diff --git a/grain/_src/python/dataset/transformations/flatmap_test.py b/grain/_src/python/dataset/transformations/flatmap_test.py index 811c8ef6f..f20b2c186 100644 --- a/grain/_src/python/dataset/transformations/flatmap_test.py +++ b/grain/_src/python/dataset/transformations/flatmap_test.py @@ -258,6 +258,104 @@ def test_random_access_flatmap_then_shuffle_keeps_nones(self): [["c0", "c1"], ["d0", "d1"], ["d2", "b0"]], ) + def test_random_access_flatmap_then_shuffle_then_slice(self): + flatmap_ds = flatmap.FlatMapMapDataset( + dataset.MapDataset.source(["a0 a1", "b0 b1", "c0 c1", "d0 d1"]), + SplitSentence(2), + ) + # The slice is applied on top of the shuffle, so it selects positions + # within the shuffled stream + # ["c0", "d0", "d1", "b1", "a0", "c1", "a1", "b0"] + # rather than within the unshuffled one. These are positions 1, 2 and 3. + sliced_ds = flatmap_ds.shuffle(seed=42)[1:4] + self.assertEqual( + [sliced_ds[i] for i in range(len(sliced_ds))], + ["d0", "d1", "b1"], + ) + + def test_random_access_flatmap_then_shuffle_then_strided_slice(self): + flatmap_ds = flatmap.FlatMapMapDataset( + dataset.MapDataset.source(["a0 a1", "b0 b1", "c0 c1", "d0 d1"]), + SplitSentence(2), + ) + # Positions 2, 4 and 6 of the shuffled stream. + sliced_ds = flatmap_ds.shuffle(seed=42)[2::2] + self.assertEqual( + [sliced_ds[i] for i in range(len(sliced_ds))], + ["d1", "a0", "a1"], + ) + + def test_random_access_flatmap_then_shuffle_then_shard(self): + # Sharding is a strided slice, so shuffling before sharding must still hand + # every word to exactly one shard. Because the shuffle permutes the + # flattened index space, the two words of a sentence can land in different + # shards: "d0" goes to shard 1 while "d1" goes to shard 0. + flatmap_ds = flatmap.FlatMapMapDataset( + dataset.MapDataset.source(["a0 a1", "b0 b1", "c0 c1", "d0 d1"]), + SplitSentence(2), + ) + shuffled_ds = flatmap_ds.shuffle(seed=42) + + shard_0 = shuffled_ds[0::2] + self.assertEqual( + [shard_0[i] for i in range(len(shard_0))], + ["c0", "d1", "a0", "a1"], + ) + + shard_1 = shuffled_ds[1::2] + self.assertEqual( + [shard_1[i] for i in range(len(shard_1))], + ["d0", "b1", "c1", "b0"], + ) + + # Together the shards cover the epoch exactly once: no word is dropped and + # none is served twice. + self.assertCountEqual( + list(shard_0) + list(shard_1), + ["a0", "a1", "b0", "b1", "c0", "c1", "d0", "d1"], + ) + + def test_random_access_slice_then_flatmap_then_shuffle(self): + # Here the slice is applied to the *sentences*, before the flat map, so the + # flattened index space only spans the sentences that survived the slice. + parent_ds = dataset.MapDataset.source( + ["z0 z1", "a0 a1", "b0 b1", "c0 c1", "d0 d1"] + )[1:] + flatmap_ds = flatmap.FlatMapMapDataset(parent_ds, SplitSentence(2)) + self.assertLen(flatmap_ds, 8) + + # Dropping the leading sentence leaves exactly the four sentences used by + # `test_random_access_flatmap_then_shuffle`, so the shuffled order is + # identical to that test: a slice below the flat map is invisible to the + # shuffle above it. + shuffled_ds = flatmap_ds.shuffle(seed=42) + self.assertEqual( + [shuffled_ds[i] for i in range(len(shuffled_ds))], + ["c0", "d0", "d1", "b1", "a0", "c1", "a1", "b0"], + ) + + def test_random_access_strided_slice_then_flatmap_then_shuffle(self): + # A strided slice of the sentences halves the flattened index space, so the + # shuffle permutes 4 words rather than 8. + parent_ds = dataset.MapDataset.source(["a0 a1", "b0 b1", "c0 c1", "d0 d1"])[ + ::2 + ] + flatmap_ds = flatmap.FlatMapMapDataset(parent_ds, SplitSentence(2)) + self.assertEqual( + [flatmap_ds[i] for i in range(len(flatmap_ds))], + ["a0", "a1", "c0", "c1"], + ) + + # Note this uses a different seed from the tests above: over a 4-element + # index space `seed=42` happens to be the identity permutation, which would + # make for a misleading test. `seed=1` is a derangement, so every word + # visibly moves and the two words of a sentence are pulled apart. + shuffled_ds = flatmap_ds.shuffle(seed=1) + self.assertEqual( + [shuffled_ds[i] for i in range(len(shuffled_ds))], + ["a1", "c0", "c1", "a0"], + ) + class Unbatch(transforms.FlatMap):