def test_reproduce_eval_data(
      self,
      per_device_batch_size,
      local_device_count,
      host_count,
      host_id,
      pad_eval,
  ):
    """Check whether loading dataset gets back original data."""

    loaded_data = data_utils.load_dataset(
        patterns=self.data_patterns[0],
        decode_fn=self.decode_fn,
        preprocess_fn=DataTest.preprocess_fn,
        collater_fn=lambda x: x,
        is_training=False,
        per_device_batch_size=per_device_batch_size,
        local_device_count=local_device_count,
        host_count=host_count,
        host_id=host_id,
        pad_eval=pad_eval,
    )
    data_list = list(loaded_data)
    recovered_data = np.array([batch['x'] for batch in data_list])

    # Compute how many non-zero elements the host should have
    host_bsz = per_device_batch_size * local_device_count
    total_bsz = host_bsz * host_count
    total_remainder = self.dataset_size % total_bsz
    remainder_per_host = total_remainder // host_count
    host_remainder = total_remainder % host_count
    n_nonzero_elements = self.dataset_size // total_bsz * host_bsz + remainder_per_host
    if host_id < host_remainder:
      n_nonzero_elements += 1
    n_batches = math.ceil(self.dataset_size / total_bsz)

    # Check data has correct shape
    self.assertEqual(
        recovered_data.shape,
        (n_batches, local_device_count, per_device_batch_size, self.data_dim))
    reshaped_data = recovered_data.reshape(-1, self.data_dim)

    # # Check if all samples in dataset are present in original data
    comparison_data = reshaped_data[:n_nonzero_elements]
    reshaped_original_data = self.data1.reshape(-1, self.data_dim)
    arr_eq = np.expand_dims(comparison_data,
                            1) == np.expand_dims(reshaped_original_data, 0)
    self.assertTrue(np.all(np.any(np.all(arr_eq, axis=-1), axis=1), axis=0))

    if pad_eval:
      # Check if nonzero sample weights correspond to expected nonzero elements
      sample_weights = np.array(
          [batch['sample_weights'] for batch in data_list]).reshape(-1)
      self.assertTrue(np.all(sample_weights[:n_nonzero_elements] == 1))
      self.assertTrue(np.all(sample_weights[n_nonzero_elements:] == 0))

      # Check if values are zero for sample weight zero
      self.assertTrue(np.all(reshaped_data[n_nonzero_elements:] == 0))
Check whether loading dataset gets back original data.
language/mentionmemory/utils/data_utils_test.py