).astype(np.float32)
        }

    dataset_path = self.create_tempfile('tfrecord').full_path
    path_map = {'serving_default': dataset_path}
    repr_dataset.TfRecordRepresentativeDatasetSaver(path_map).save(
        {'serving_default': data_gen()}
    )

    config = qc.QuantizationConfig(
        static_range_ptq_preset=qc.StaticRangePtqPreset(
            representative_datasets=[

    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(

    """
    raise NotImplementedError('Method "save" is not implemented.')


@tf_export.tf_export(
    'quantization.experimental.TfRecordRepresentativeDatasetSaver'
)
class TfRecordRepresentativeDatasetSaver(RepresentativeDatasetSaver):
  """Representative dataset saver in TFRecord format.

  Saves representative datasets for quantization calibration in TFRecord format.

    _, path_map[signature_key] = tempfile.mkstemp(
        suffix='.tfrecord', prefix=signature_key
    )
    expected_input_key_map[signature_key] = signature_def.inputs.keys()

  return repr_dataset.TfRecordRepresentativeDatasetSaver(
      path_map=path_map,
      expected_input_key_map=expected_input_key_map,
  ).save(representative_dataset_map)


def _run_static_range_qat(
    src_saved_model_path: str,

    representative_dataset: repr_dataset.RepresentativeDataset = [
        {'input_tensor': rng.uniform(size=(1, 1024)).astype(np.float32)}
        for _ in range(4)
    ]
    dataset_file_map = repr_dataset.TfRecordRepresentativeDatasetSaver(
        {'serving_default': os.path.join(self._input_saved_model_path, 'repr')}
    ).save({'serving_default': representative_dataset})

    tags = {tag_constants.SERVING}