From 03e58923dc91dc6c086f67dc817b75e2fec47505 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Thu, 24 Sep 2026 16:57:29 -0700 Subject: [PATCH] internal PiperOrigin-RevId: 987853082 --- .../python/dataset/transformations/shuffle_test.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/grain/_src/python/dataset/transformations/shuffle_test.py b/grain/_src/python/dataset/transformations/shuffle_test.py index 4524d2136..9b0820b45 100644 --- a/grain/_src/python/dataset/transformations/shuffle_test.py +++ b/grain/_src/python/dataset/transformations/shuffle_test.py @@ -118,6 +118,17 @@ def test_slice_then_shuffle_by_index(self): self.assertEqual(shuffled_ds[2], "e") self.assertEqual(shuffled_ds[3], "b") + def test_slice_then_shuffle_twice(self): + # Slice first to restrict elements to ["b", "c", "d", "e"], then shuffle + # the 4-element sliced dataset. + ds = dataset.MapDataset.source(["a", "b", "c", "d", "e", "f"])[1:5] + self.assertEqual(list(ds), ["b", "c", "d", "e"]) + shuffled_ds = ds.shuffle(seed=1) + shuffled_ds = shuffled_ds.shuffle(seed=10) + self.assertLen(shuffled_ds, 4) + self.assertEqual(shuffled_ds[0], "e") + self.assertEqual(shuffled_ds[3], "b") + class WindowShuffleMapDatasetTest(absltest.TestCase):