# Mapping of signature def key -> SignatureDef.
_SignatureDefMap = Mapping[str, meta_graph_pb2.SignatureDef]


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:

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::PyFunctionLibrary;

// Returns signature key -> `SignatureDef` mapping, excluding the signature for
// initialization op, which is only used during initialization.
// TODO: b/314124142 - Remove the need for this function.
absl::flat_hash_map<std::string, SignatureDef> GetSignatureDefMapFromBundle(

#include "tensorflow/core/protobuf/saver.pb.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {

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

StaticRangePtqComponent::StaticRangePtqComponent(
    absl::Nonnull<MLIRContext*> ctx,

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

    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,
    const tensorflow::quantization::PyFunctionLibrary& py_function_library);

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

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

  // Runs the static-range post-training quantization (PTQ) on `module_op`.
  absl::StatusOr<ModuleOp> Run(

#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,
  // )