Use etree instead of a custom module for parallel_map.

PiperOrigin-RevId: 618901084

Author: marcenacp

tensorflow_datasets/core/dataset_builder.py

@@ -54,11 +54,12 @@
 from tensorflow_datasets.core.utils import file_utils
 from tensorflow_datasets.core.utils import gcs_utils
 from tensorflow_datasets.core.utils import read_config as read_config_lib
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import apache_beam as beam
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 import termcolor
+import tree
+
 
 ListOrTreeOrElem = type_utils.ListOrTreeOrElem
 Tree = type_utils.Tree
@@ -794,7 +795,7 @@ def build_single_data_source(
             f" {' or '.join(args)}?"
         )
 
-    all_ds = tree_utils.map_structure(build_single_data_source, split)
+    all_ds = tree.map_structure(build_single_data_source, split)
     return all_ds
 
   @tfds_logging.as_dataset()
@@ -908,7 +909,7 @@ def as_dataset(
         read_config=read_config,
         as_supervised=as_supervised,
     )
-    all_ds = tree_utils.map_structure(build_single_dataset, split)
+    all_ds = tree.map_structure(build_single_dataset, split)
     return all_ds
 
   def _build_single_dataset(
@@ -961,7 +962,7 @@ def lookup_nest(features: Dict[str, Any]) -> Tuple[Any, ...]:
         Returns:
           A tuple with elements structured according to `supervised_keys`
         """
-        return tree_utils.map_structure(
+        return tree.map_structure(
             lambda key: features[key], self.info.supervised_keys
         )
 

tensorflow_datasets/core/dataset_metadata.py

@@ -19,8 +19,8 @@
 import functools
 
 from etils import epath
+from etils import etree
 from tensorflow_datasets.core import constants
-from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.utils import resource_utils
 
 
@@ -68,13 +68,11 @@ def valid_tags() -> list[str]:
 
 def valid_tags_with_comments() -> str:
   """Returns valid tags (one per line) with comments."""
-  return "\n".join(
-      [
-          line
-          for line in _get_valid_tags_text().split("\n")
-          if not line.startswith("#")
-      ]
-  )
+  return "\n".join([
+      line
+      for line in _get_valid_tags_text().split("\n")
+      if not line.startswith("#")
+  ])
 
 
 @functools.lru_cache(maxsize=256)
@@ -103,6 +101,4 @@ def _read_files(path: epath.Path) -> dict[str, str]:
   for inode in path.iterdir():
     if inode.name in _METADATA_FILES:
       name2path[inode.name] = path.joinpath(inode.name)
-  return utils.tree.parallel_map(
-      lambda f: f.read_text(encoding="utf-8"), name2path
-  )
+  return etree.parallel_map(lambda f: f.read_text(encoding="utf-8"), name2path)

tensorflow_datasets/core/dataset_utils.py

@@ -27,9 +27,9 @@
 import numpy as np
 from tensorflow_datasets.core import logging as tfds_logging
 from tensorflow_datasets.core import utils
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+import tree
 
 Tree = type_utils.Tree
 Tensor = type_utils.Tensor
@@ -80,7 +80,7 @@ def element_spec(self) -> Tree[enp.ArraySpec]:
 
 def _eager_dataset_iterator(ds: tf.data.Dataset) -> Iterator[NumpyElem]:
   for elem in ds:
-    yield tree_utils.map_structure(_elem_to_numpy_eager, elem)
+    yield tree.map_structure(_elem_to_numpy_eager, elem)
 
 
 def _graph_dataset_iterator(ds_iter, graph: tf.Graph) -> Iterator[NumpyElem]:
@@ -105,7 +105,7 @@ def _assert_ds_types(nested_ds: Tree[TensorflowElem]) -> None:
         isinstance(el, (tf.Tensor, tf.RaggedTensor))
         or isinstance(el, tf.data.Dataset)
     ):
-      nested_types = tree_utils.map_structure(type, nested_ds)
+      nested_types = tree.map_structure(type, nested_ds)
       raise TypeError(
           'Arguments to as_numpy must be tf.Tensors or tf.data.Datasets. '
           f'Got: {nested_types}.'
@@ -197,7 +197,7 @@ def as_numpy(dataset: Tree[TensorflowElem]) -> Tree[NumpyElem]:
   """
   _assert_ds_types(dataset)
   if tf.executing_eagerly():
-    return tree_utils.map_structure(_elem_to_numpy_eager, dataset)
+    return tree.map_structure(_elem_to_numpy_eager, dataset)
   else:
     return _nested_to_numpy_graph(dataset)
 

tensorflow_datasets/core/decode/base.py

@@ -18,8 +18,8 @@
 import abc
 import functools
 
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+import tree
 
 
 class Decoder(abc.ABC):
@@ -69,7 +69,7 @@ def setup(self, *, feature):
   def dtype(self):
     """Returns the `dtype` after decoding."""
     tensor_info = self.feature.get_tensor_info()
-    return tree_utils.map_structure(lambda t: t.dtype, tensor_info)
+    return tree.map_structure(lambda t: t.dtype, tensor_info)
 
   @abc.abstractmethod
   def decode_example(self, serialized_example):
@@ -135,7 +135,7 @@ class SkipDecoding(Decoder):
   @property
   def dtype(self):
     tensor_info = self.feature.get_serialized_info()
-    return tree_utils.map_structure(lambda t: t.dtype, tensor_info)
+    return tree.map_structure(lambda t: t.dtype, tensor_info)
 
   def decode_example(self, serialized_example):
     """Forward the serialized feature field."""

tensorflow_datasets/core/download/download_manager.py

@@ -35,8 +35,8 @@
 from tensorflow_datasets.core.download import resource as resource_lib
 from tensorflow_datasets.core.download import util
 from tensorflow_datasets.core.utils import shard_utils
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
+import tree
 
 
 # pylint: disable=logging-fstring-interpolation
@@ -825,10 +825,6 @@ def _read_url_info(url_path: epath.PathLike) -> checksums.UrlInfo:
 
 def _map_promise(map_fn, all_inputs):
   """Map the function into each element and resolve the promise."""
-  all_promises = tree_utils.map_structure(
-      map_fn, all_inputs
-  )  # Apply the function
-  res = tree_utils.map_structure(
-      lambda p: p.get(), all_promises
-  )  # Wait promises
+  all_promises = tree.map_structure(map_fn, all_inputs)  # Apply the function
+  res = tree.map_structure(lambda p: p.get(), all_promises)  # Wait promises
   return res

tensorflow_datasets/core/download/download_manager_test.py

@@ -25,16 +25,14 @@
 from absl.testing import parameterized
 from etils import epath
 import promise
-
 import tensorflow as tf
-
 from tensorflow_datasets import testing
 from tensorflow_datasets.core.download import checksums as checksums_lib
 from tensorflow_datasets.core.download import download_manager as dm
 from tensorflow_datasets.core.download import downloader
 from tensorflow_datasets.core.download import extractor
 from tensorflow_datasets.core.download import resource as resource_lib
-from tensorflow_datasets.core.utils import tree_utils
+import tree
 
 ZIP = resource_lib.ExtractMethod.ZIP
 TAR = resource_lib.ExtractMethod.TAR
@@ -46,7 +44,7 @@ def _sha256(str_):
 
 
 def _as_path(nested_paths):
-  return tree_utils.map_structure(epath.Path, nested_paths)
+  return tree.map_structure(epath.Path, nested_paths)
 
 
 def _info_path(path):

tensorflow_datasets/core/features/dataset_feature.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 """Dataset feature for nested datasets."""
+
 from __future__ import annotations
 
 import dataclasses
@@ -26,9 +27,9 @@
 from tensorflow_datasets.core.features import tensor_feature
 from tensorflow_datasets.core.features import top_level_feature
 from tensorflow_datasets.core.utils import py_utils
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+import tree
 
 
 @dataclasses.dataclass(frozen=True)
@@ -120,7 +121,7 @@ def get_tensor_info(self):
     """Shape of one element of the dataset."""
     # Add the dataset level
     tensor_info = self._feature.get_tensor_info()
-    return tree_utils.map_structure(_add_dataset_lvl, tensor_info)
+    return tree.map_structure(_add_dataset_lvl, tensor_info)
 
   def get_tensor_spec(self) -> tf.data.DatasetSpec:
     return tf.data.DatasetSpec(element_spec=self._feature.get_tensor_spec())
@@ -129,7 +130,7 @@ def get_tensor_spec(self) -> tf.data.DatasetSpec:
   def get_serialized_info(self):
     # Add the dataset level and the number of elements in the dataset
     tensor_info = super().get_serialized_info()
-    return tree_utils.map_structure(_add_dataset_lvl, tensor_info)
+    return tree.map_structure(_add_dataset_lvl, tensor_info)
 
   def encode_example(
       self,
@@ -146,7 +147,7 @@ def encode_example(
 
     # Empty datasets return empty arrays
     if not ds_elements:
-      return tree_utils.map_structure(
+      return tree.map_structure(
           sequence_feature.build_empty_np, self.get_serialized_info()
       )
 

tensorflow_datasets/core/features/feature.py

@@ -38,10 +38,10 @@
 from tensorflow_datasets.core.utils import np_utils
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import tf_utils
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 from tensorflow_datasets.core.utils.lazy_imports_utils import tf_agents
+import tree
 
 from google.protobuf import descriptor
 from google.protobuf import json_format
@@ -364,14 +364,14 @@ def get_tensor_spec(self) -> TreeDict[tf.TensorSpec]:
     of the dataset. For example, currently this method does not support
     RaggedTensorSpec.
     """
-    return tree_utils.map_structure(
+    return tree.map_structure(
         lambda ti: ti.to_tensor_spec(), self.get_tensor_info()
     )
 
   @functools.cached_property
   def shape(self):
     """Return the shape (or dict of shape) of this FeatureConnector."""
-    return tree_utils.map_structure(lambda t: t.shape, self.get_tensor_info())
+    return tree.map_structure(lambda t: t.shape, self.get_tensor_info())
 
   @functools.cached_property
   def dtype(self) -> TreeDict[tf.dtypes.DType]:
@@ -381,9 +381,7 @@ def dtype(self) -> TreeDict[tf.dtypes.DType]:
 
   @functools.cached_property
   def np_dtype(self) -> TreeDict[np.dtype]:
-    return tree_utils.map_structure(
-        lambda t: t.np_dtype, self.get_tensor_info()
-    )
+    return tree.map_structure(lambda t: t.np_dtype, self.get_tensor_info())
 
   # For backwards compatibility: now it is named np_dtype.
   @functools.cached_property
@@ -397,7 +395,7 @@ def convert_to_tensorflow(value):
         return tf.dtypes.as_dtype(value)
       return value.tf_dtype
 
-    return tree_utils.map_structure(convert_to_tensorflow, self.np_dtype)
+    return tree.map_structure(convert_to_tensorflow, self.np_dtype)
 
   @classmethod
   def cls_from_name(cls, python_class_name: str) -> Type['FeatureConnector']:

tensorflow_datasets/core/features/sequence_feature.py

@@ -27,8 +27,8 @@
 from tensorflow_datasets.core.features import top_level_feature
 from tensorflow_datasets.core.proto import feature_pb2
 from tensorflow_datasets.core.utils import py_utils
-from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
+import tree
 
 Json = type_utils.Json
 
@@ -121,14 +121,14 @@ def get_tensor_info(self):
     """See base class for details."""
     # Add the additional length dimension to every shape
     tensor_info = self._feature.get_tensor_info()
-    return tree_utils.map_structure(self._add_length_dim, tensor_info)
+    return tree.map_structure(self._add_length_dim, tensor_info)
 
   @py_utils.memoize()
   def get_serialized_info(self):
     """See base class for details."""
     # Add the additional length dimension to every serialized features
     tensor_info = self._feature.get_serialized_info()
-    return tree_utils.map_structure(self._add_length_dim, tensor_info)
+    return tree.map_structure(self._add_length_dim, tensor_info)
 
   def encode_example(self, example_dict):
     # Convert nested dict[list] into list[nested dict]
@@ -143,9 +143,7 @@ def encode_example(self, example_dict):
 
     # Empty sequences return empty arrays
     if not sequence_elements:
-      return tree_utils.map_structure(
-          build_empty_np, self.get_serialized_info()
-      )
+      return tree.map_structure(build_empty_np, self.get_serialized_info())
 
     # Encode each individual elements
     sequence_elements = [

tensorflow_datasets/core/folder_dataset/compute_split_utils.py

@@ -25,11 +25,11 @@
 from typing import Optional, Type, Union, cast
 
 from etils import epath
+from etils import etree
 from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core import lazy_imports_lib
 from tensorflow_datasets.core import naming
 from tensorflow_datasets.core import splits as split_lib
-from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.proto import dataset_info_pb2
 from tensorflow_datasets.core.utils.lazy_imports_utils import apache_beam as beam
 from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_data_source
@@ -230,14 +230,14 @@ def _compute_split_statistics(
   # Compute all shard info in parallel
   split_to_shard_infos = cast(
       Mapping[str, Sequence[_ShardInfo]],
-      utils.tree.parallel_map(
+      etree.parallel_map(
           functools.partial(
               _process_shard,
               data_dir=filename_template.data_dir,
               adapter=adapter,
           ),
           split_files,
-          report_progress=True,
+          progress_bar=True,
       ),
   )
   # Create the SplitInfo for all splits