dataset at a file system or when you need to pass the representative dataset
  as files.
  """

  def save(
      self, representative_dataset: RepresentativeDatasetMapping
  ) -> Mapping[str, _RepresentativeDatasetFile]:
    """Saves the representative dataset.

    Args:
      representative_dataset: RepresentativeDatasetMapping which is a
        signature_def_key -> representative dataset mapping.
    """

limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h"

#include <string>
#include <utility>

#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"

Dan Suh <******@****.***> 1706503553 -0800

  * If `len(signature_keys) > 1`, then `representative_dataset` should be a
    mapping where the keys exactly match the elements in `signature_keys`.
  * If `len(signature_keys) == 1`, then both a mapping and a single instance of
    `RepresentativeDataset` are allowed.
  * This function also assumes `len(signature_keys) > 0`.

  Args:
    representative_dataset: A `RepresentativeDataset` or a map of string to

  )

  if (
      representative_dataset is not None
      and quantization_options.representative_datasets
  ):
    raise ValueError(
        'Do not specify both the `representative_dataset` argument and'
        ' the `representative_datasets` field in `QuantizationOptions`.'
    )

  saved_representative_dataset = quantization_options.representative_datasets
  if representative_dataset is not None:

        "@llvm-project//mlir:Support",
        "@local_tsl//tsl/platform:errors",
        "@local_tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "representative_dataset",
    srcs = ["representative_dataset.cc"],
    hdrs = ["representative_dataset.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",

        representative_datasets=dataset_file_map,
    )

    with self.assertRaisesRegex(
        ValueError,
        'Do not specify both the `representative_dataset` argument and the'
        ' `representative_datasets` field in `QuantizationOptions`',
    ):
      quantize_model.quantize(
          self._input_saved_model_path,
          self._output_saved_model_path,

# limitations under the License.
# ==============================================================================
"""Tests for representative_dataset.py."""
import random

import numpy as np

from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as repr_dataset
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.client import session

        "//tensorflow/python/saved_model:tag_constants",
        "//tensorflow/python/trackable:autotrackable",
    ],
)

pytype_strict_library(
    name = "representative_dataset",
    srcs = [
        "representative_dataset.py",
    ],
    visibility = ["//visibility:public"],
    deps = [
        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_py",

from typing import Any

from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as rd

# LINT.IfChange(static_range_ptq)
def static_range_ptq(
    src_saved_model_path: str,
    dst_saved_model_path: str,
    quantization_config_serialized: bytes,
    *,
    signature_keys: list[str],