PYBIND11_OVERRIDE_PURE(std::optional<MinMaxValue>, PyFunctionLibrary,
                           get_calibration_min_max_value,
                           calibration_statistics, calibration_options);
  }
};

}  // namespace

PYBIND11_MODULE(pywrap_function_lib, m) {
  py::class_<PyFunctionLibrary, PyFunctionLibraryTrampoline>(
      m, "PyFunctionLibrary")
      .def(py::init<>())

    py_function_library: py_function_lib.PyFunctionLibrary,
) -> Any: ...  # Status

# LINT.ThenChange()

# LINT.IfChange(quantize_weight_only)
def quantize_weight_only(
    src_saved_model_path: str,
    dst_saved_model_path: str,
    quantization_options_serialized: bytes,
    *,
    signature_def_map_serialized: dict[str, bytes],
    py_function_library: py_function_lib.PyFunctionLibrary,
) -> Any: ...  # Status

    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,
) -> Any: ...  # Status

# LINT.ThenChange()

# LINT.IfChange(weight_only_ptq)
def weight_only_ptq(
    src_saved_model_path: str,
    dst_saved_model_path: str,

#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace py = pybind11;

namespace {

using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::QuantizationOptions;
using ::tensorflow::quantization::QuantizeDynamicRangePtq;
using ::tensorflow::quantization::QuantizeQatModel;
using ::tensorflow::quantization::QuantizeStaticRangePtq;

using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;

StaticRangePtqComponent::StaticRangePtqComponent(
    absl::Nonnull<MLIRContext*> ctx,
    absl::Nonnull<const PyFunctionLibrary*> py_function_library,
    const absl::string_view src_saved_model_path,
    std::vector<std::string> signature_keys,

          quant_opts_serialized,
          signature_keys=signature_def_keys,
          signature_def_map_serialized={},
          function_aliases={},
          py_function_library=py_function_lib.PyFunctionLibrary(),
          representative_dataset_file_map_serialized=None,
      )

  def test_quantize_model_fails_when_invalid_quant_options_type(self):
    src_saved_model_path = self.create_tempdir().full_path

namespace stablehlo::quantization::pywrap {

using ::mlir::quant::stablehlo::QuantizeStaticRangePtq;
using ::mlir::quant::stablehlo::QuantizeWeightOnlyPtq;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::PyFunctionLibrary;

// Note for maintainers: the definitions should ONLY mirror existing functions
// defined in different targets. Do not include any extra business logic that
// causes divergence from the semantics of mirrored functions.

    const std::vector<std::string>& signature_keys,
    const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
        signature_def_map,
    const tensorflow::quantization::PyFunctionLibrary& py_function_library);

// Function used by the pywrap_quantization module to mirror
// `::mlir::quant::stablehlo::QuantizeWeightOnlyPtq`.
absl::Status PywrapQuantizeWeightOnlyPtq(

        signature_keys=list(signature_def_map.keys()),
        signature_def_map_serialized=signature_def_map_serialized,
        py_function_library=py_function_lib.PyFunctionLibrary(),
    )
  elif _has_quantization_method(config.specs, 'weight_only_ptq'):
    pywrap_quantization.weight_only_ptq(
        src_saved_model_path,
        dst_saved_model_path,

using ::stablehlo::quantization::GetReportFilePath;
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;

WeightOnlyPtqComponent::WeightOnlyPtqComponent(absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK