).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=[

    _, 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}