#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {

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

// weight-only quantization.
//
// `signature_keys` specify the signatures that correspond to functions to be
// quantized. `signature_def_map` connects the signature keys to
// `SignatureDef`s.
//
// Returns a non-OK status when the quantization is not successful.
// LINT.IfChange
absl::Status QuantizeWeightOnlyPtq(
    absl::string_view src_saved_model_path,
    absl::string_view dst_saved_model_path,

        `dst_saved_model_path`.

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

        Raises `StatusNotOk` exception if when the run was unsuccessful.
        )pbdoc",

      const ExportedModel& exported_model,
      absl::string_view src_saved_model_path,
      const std::unordered_set<std::string>& tags,
      const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
          signature_def_map) const = 0;
  // LINT.ThenChange(
  //     pywrap_function_lib.pyi:save_exported_model,
  //     py_function_lib.py:save_exported_model,
  // )

    Args:
      path_map: Signature def key -> path mapping. Each path is a TFRecord file
        to which a `RepresentativeDataset` is saved. The signature def keys
        should be a subset of the `SignatureDef` keys of the
        `representative_dataset` argument of the `save()` call.
      expected_input_key_map: Signature def key -> expected input keys. If set,

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

namespace py = ::pybind11;

namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::MinMaxValue;
using ::tensorflow::SignatureDef;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RepresentativeDatasetFile;

    : public internal::SerializedProtobufCaster<
          ::stablehlo::quantization::CalibrationOptions> {};

template <>
struct type_caster<tensorflow::SignatureDef>
    : public internal::SerializedProtobufCaster<tensorflow::SignatureDef> {};

template <>
struct type_caster<tensorflow::GraphDef>
    : public internal::SerializedProtobufCaster<tensorflow::GraphDef> {};

template <>

def _serialize_signature_def_map(
    signature_def_map: _SignatureDefMap,
) -> dict[str, bytes]:
  """Serializes SignatureDef values in `signature_def_map`.

  Args:
    signature_def_map: Signature key -> SignatureDef mapping.

  Returns:
    Signature def map where the values (`SignatureDef`) are serialized.
  """
  signature_def_map_serialized = {}
  for key, signature_def in signature_def_map.items():

      absl::string_view src_saved_model_path,
      absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases,
      std::unordered_set<std::string> tags,
      absl::flat_hash_map<std::string, tensorflow::SignatureDef>
          signature_def_map,
      std::vector<std::string> signature_keys);

  // Runs calibration on `module_op` and returns a calibrated ModuleOp with
  // calibrated statistics embedded.

using ::stablehlo::quantization::io::CreateTmpDir;
using ::stablehlo::quantization::io::GetLocalTmpFileName;
using ::stablehlo::quantization::io::ListDirectory;
using ::tensorflow::AssetFileDef;
using ::tensorflow::SignatureDef;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;