def test_training_data_deterministic(self, is_training=True):
    """Make sure loading data is deterministic given seed."""
    data_arrays = []
    for _ in range(2):
      loaded_dataset = data_utils.load_dataset(
          patterns=self.data_patterns,
          decode_fn=self.decode_fn,
          preprocess_fn=DataTest.preprocess_fn,
          collater_fn=lambda x: x,
          is_training=is_training,
          per_device_batch_size=4,
          local_device_count=1,
          host_count=2,
          host_id=0,
          seed=0,
      )
      data_iterator = iter(loaded_dataset)
      data_array = np.array([
          data_iterator.get_next()['x'] for _ in range(self.dataset_size // 4)
      ])
      data_arrays.append(data_array)

    self.assertTrue(np.array_equal(*data_arrays))
Make sure loading data is deterministic given seed.
language/mentionmemory/utils/data_utils_test.py