Migrate tf.data.experimental.RandomDataset to `tf.data.Dataset.rand…

…om`. PiperOrigin-RevId: 731853332
google-research · Feb 27, 2025 · 2c2c0d5 · 2c2c0d5
1 parent fce4b1a
commit 2c2c0d5
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/t5/data/preprocessors_test.py b/t5/data/preprocessors_test.py
@@ -2440,9 +2440,7 @@ def test_cast_seeds_works_with_tf_rngs(self):
     base_seed = 42
     random_ds_seeds = np.arange(base_seed, base_seed + 2).reshape(-1, 2)
     random_ds_seeds = tuple(tuple(s) for s in random_ds_seeds)
-    seed_ds = tf.nest.map_structure(
-        tf.data.experimental.RandomDataset, random_ds_seeds
-    )
+    seed_ds = tf.nest.map_structure(tf.data.Dataset.random, random_ds_seeds)
     range_ds = tf.data.Dataset.from_tensor_slices(range(1))
     zip_ds = tf.data.Dataset.zip(range_ds, seed_ds)
     for d in zip_ds: