py_function_library.SaveExportedModel(
            dst_saved_model_path, *exported_model, src_saved_model_path, tags,
            signature_def_map);

        return absl::OkStatus();
      },
      R"pbdoc(
      Quantizes a model that went through quantization-aware training (QAT)
      saved at `src_saved_model_path`. The resulting model will be saved to
      `dst_saved_model_path`. Returns an OK sataus when successful, otherwise

    dst_saved_model_path = self.create_tempdir().full_path
    signature_def_keys = ['serving_default']
    invalid_quant_opts_object = ('a', 'b', 'c')

    with self.assertRaisesRegex(TypeError, 'incompatible function arguments'):
      pywrap_quantize_model.quantize_ptq_static_range(
          src_saved_model_path,
          dst_saved_model_path,
          invalid_quant_opts_object,

# TODO: b/310594193 - Export API to pip package.
def quantize_saved_model(
    src_saved_model_path: str,
    dst_saved_model_path: str,
    config: qc.QuantizationConfig,
) -> None:
  """Quantizes a saved model.

  Args:
    src_saved_model_path: Path to the directory for the source SavedModel.
    dst_saved_model_path: Path to the directory for the destination SavedModel.
    config: Quantization configuration.

  Raises:

        R"pbdoc(
        Runs static-range post-training quantization (PTQ) on a SavedModel at
        `src_saved_model_path` and saves the resulting model to
        `dst_saved_model_path`.

        The user should pass a serialized `QuantizationConfig` for the
        `quantization_config_serialized` argument, and a signature key ->

class PyFunctionLibrary {
 public:
  virtual ~PyFunctionLibrary() = default;

  // Saves `exported_model` to `dst_saved_model_path` as SavedModel.
  // `src_saved_model_path` is the path to the source SavedModel from which the
  // exported model is produced. It is used to copy the asset files to
  // `dst_saved_model_path`. `tags` will be attached to the saved
  // `MetaGraphDef`. `signature_def_map` will be passed to the

// Function used by the pywrap_quantization module to mirror
// `::mlir::quant::stablehlo::QuantizeStaticRangePtq`.
absl::Status PywrapQuantizeStaticRangePtq(
    absl::string_view src_saved_model_path,
    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
    const std::vector<std::string>& signature_keys,
    const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
        signature_def_map,

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

# LINT.IfChange(quantize_qat_model)
def quantize_qat_model(
    src_saved_model_path: str,
    dst_saved_model_path: str,
    quantization_options_serialized: bytes,
    *,
    signature_keys: list[str],
    signature_def_map_serialized: dict[str, bytes],
    py_function_library: py_function_lib.PyFunctionLibrary,

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],
    signature_def_map_serialized: dict[str, bytes],
    py_function_library: py_function_lib.PyFunctionLibrary,

    absl::string_view src_saved_model_path,
    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
    const std::vector<std::string>& signature_keys,
    const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
    const PyFunctionLibrary& py_function_library) {
  return QuantizeStaticRangePtq(src_saved_model_path, dst_saved_model_path,
                                config, signature_keys, signature_def_map,

        src_saved_model_path,
        dst_saved_model_path,
        quantization_options,
        signature_def_map,
    )
  else:
    _run_static_range_ptq(
        src_saved_model_path,
        dst_saved_model_path,
        quantization_options,
        saved_representative_dataset,
        signature_def_map,
    )

  return saved_model_load.load(dst_saved_model_path)


def _dynamic_range_quantize(