num_samples = 2

    def data_gen():
      for _ in range(num_samples):
        yield {'x': [1, 2]}

    repr_ds_map = {'serving_default': data_gen()}
    saver = repr_dataset.TfRecordRepresentativeDatasetSaver(path_map)
    dataset_file_map = saver.save(repr_ds_map)
    self.assertCountEqual(dataset_file_map.keys(), ['serving_default'])

    dataset_map = repr_dataset.TfRecordRepresentativeDatasetLoader(

    for signature_def_key in repr_ds_map:
      # Replaces the dataset with a new dataset where tf.Tensors are replaced
      # by their evaluated values.
      ds = repr_ds_map[signature_def_key]
      repr_ds_map[signature_def_key] = rd.replace_tensors_by_numpy_ndarrays(
          ds, sess
      )


def _create_sample_validator(
    expected_input_keys: Collection[str],