diff --git a/grain/_src/python/dataset/transformations/interleave_test.py b/grain/_src/python/dataset/transformations/interleave_test.py index c5d647307..3cd97052e 100644 --- a/grain/_src/python/dataset/transformations/interleave_test.py +++ b/grain/_src/python/dataset/transformations/interleave_test.py @@ -291,6 +291,14 @@ def test_set_next_index_with_multiple_datasets(self): ): dataset.set_next_index(ds_iter, 0) + def test_on_par_with_cpp_backend(self): + ds1 = dataset.MapDataset.range(1).to_iter_dataset() + ds2 = dataset.MapDataset.range(10, 12).to_iter_dataset() + ds3 = dataset.MapDataset.range(20, 21).to_iter_dataset() + + ds = interleave.InterleaveIterDataset([ds1, ds2, ds3], cycle_length=2) + self.assertEqual(list(ds), [0, 10, 11, 20]) + if __name__ == "__main__": absltest.main()