Internal change

PiperOrigin-RevId: 390551401

Author: andsteing

CHANGELOG.md

@@ -85,3 +85,7 @@ Documentation:
 -   Makes `PreprocessFn` addable.
 -   Allow n-dimensional arrays (and masks) to be passed to Metrics.Average().
 -   Support slicing `PreprocessFn`.
+
+## v0.0.6
+
+-   Makes `deterministic_data` work with `tfds>4.4.0` and `tfds<=4.4.0`.

clu/deterministic_data.py

@@ -62,6 +62,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+from packaging import version
 import tensorflow as tf
 import tensorflow_datasets as tfds
 import typing_extensions
@@ -73,6 +74,9 @@
 
 AUTOTUNE = tf.data.experimental.AUTOTUNE
 
+_use_split_info = version.parse("4.4.0") < version.parse(
+    tfds.version.__version__)
+
 
 class DatasetBuilder(typing_extensions.Protocol):
   """Protocol for dataset builders (subset of tfds.core.DatasetBuilder)."""
@@ -106,15 +110,18 @@ class RemainderOptions(enum.Enum):
 def _shard_read_instruction(
     absolute_instruction,
     *,
-    split_infos: Dict[str, tfds.core.SplitInfo],
+    split_infos: Dict[str, Union[int, tfds.core.SplitInfo]],
     host_id: int,
     host_count: int,
     remainder_options: RemainderOptions,
 ) -> tfds.core.ReadInstruction:
   """Shards a single ReadInstruction. See get_read_instruction_for_host()."""
   start = absolute_instruction.from_ or 0
-  end = absolute_instruction.to or (
-      split_infos[absolute_instruction.splitname].num_examples)
+  if _use_split_info:
+    end = absolute_instruction.to or (
+        split_infos[absolute_instruction.splitname].num_examples)  # pytype: disable=attribute-error
+  else:
+    end = absolute_instruction.to or split_infos[absolute_instruction.splitname]
   assert end >= start, f"start={start}, end={end}"
   num_examples = end - start
 
@@ -208,16 +215,23 @@ def get_read_instruction_for_host(
         f"Invalid combination of host_id ({host_id}) and host_count "
         f"({host_count}).")
 
-  if dataset_info is None:
-    split_infos = {
-        split: tfds.core.SplitInfo(
-            name=split,
-            shard_lengths=[num_examples],
-            num_bytes=0,
-        ),
-    }
+  if _use_split_info:
+    if dataset_info is None:
+      split_infos = {
+          split: tfds.core.SplitInfo(
+              name=split,
+              shard_lengths=[num_examples],
+              num_bytes=0,
+          ),
+      }
+    else:
+      split_infos = dataset_info.splits
   else:
-    split_infos = dataset_info.splits
+    if dataset_info is None:
+      split_infos = {split: num_examples}
+    else:
+      split_infos = {k: v.num_examples for k, v in dataset_info.splits.items()}
+
   read_instruction = tfds.core.ReadInstruction.from_spec(split)
   sharded_read_instructions = []
   for ri in read_instruction.to_absolute(split_infos):

clu/deterministic_data_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Unit tests for the deterministic_data module."""
+import dataclasses
 import itertools
 import math
 
@@ -21,11 +22,14 @@
 
 from absl.testing import parameterized
 from clu import deterministic_data
-import dataclasses
 import jax
+from packaging import version
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
+_use_split_info = version.parse("4.4.0") < version.parse(
+    tfds.version.__version__)
+
 
 @dataclasses.dataclass
 class MyDatasetBuilder:
@@ -35,11 +39,14 @@ class MyDatasetBuilder:
   def as_dataset(self, split: tfds.core.ReadInstruction, shuffle_files: bool,
                  read_config: tfds.ReadConfig, decoders) -> tf.data.Dataset:
     del shuffle_files, read_config, decoders
-    split_infos = {
-        k: tfds.core.SplitInfo(name=k, shard_lengths=[v], num_bytes=0)
-        for k, v in self.name2len.items()
-    }
-    instructions = split.to_absolute(split_infos)
+    if _use_split_info:
+      split_infos = {
+          k: tfds.core.SplitInfo(name=k, shard_lengths=[v], num_bytes=0)
+          for k, v in self.name2len.items()
+      }
+      instructions = split.to_absolute(split_infos)
+    else:
+      instructions = split.to_absolute(self.name2len)
     assert len(instructions) == 1
     from_ = instructions[0].from_ or 0
     to = instructions[0].to or self.name2len[instructions[0].splitname]
@@ -88,12 +95,15 @@ def test_get_read_instruction_for_host_deprecated(self, num_examples: int,
         host_id=host_id,
         host_count=host_count,
         drop_remainder=drop_remainder)
-    split_infos = {
-        "test": tfds.core.SplitInfo(
-            name="test",
-            shard_lengths=[9],
-            num_bytes=0,
-        )}
+    if _use_split_info:
+      split_infos = {
+          "test": tfds.core.SplitInfo(
+              name="test",
+              shard_lengths=[9],
+              num_bytes=0,
+          )}
+    else:
+      split_infos = {"test": 9}
     self.assertEqual(
         expected.to_absolute(split_infos), actual.to_absolute(split_infos))
 

setup.py

@@ -33,7 +33,7 @@
 
 setup(
     name="clu",
-    version="0.0.5",
+    version="0.0.6",
     description=("Set of libraries for ML training loops in JAX."),
     author="Common Loop Utils Authors",
     author_email="no-reply@google.com",
@@ -51,6 +51,7 @@
         "jaxlib",
         "ml_collections",
         "numpy",
+        "packaging",
         "tensorflow",
         "tensorflow_datasets",
     ],