Switch pmap to jit and [automatic parallelization](https://jax.readth…

…edocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) PiperOrigin-RevId: 561669843
tensorflow · Sep 4, 2023 · 9ce5057 · 9ce5057
1 parent 2c66051
commit 9ce5057
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_datasets/core/dataset_utils.py b/tensorflow_datasets/core/dataset_utils.py
@@ -23,7 +23,7 @@
 from typing import Any, Callable, Iterable, Iterator, Union
 
 from etils import enp
-from etils.etree import nest as etree
+from etils.etree import nest as etree  # pylint: disable=g-importing-member
 import numpy as np
 from tensorflow_datasets.core import logging as tfds_logging
 from tensorflow_datasets.core import utils