# 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:

    EXPECT_EQ(actual.ShortDebugString(), expected.ShortDebugString());
  }

  // Builds a SignatureDef with the provided `inputs` and `outputs`.
  SignatureDef BuildSignatureDef(const std::unordered_set<string>& inputs,
                                 const std::unordered_set<string>& outputs) {
    SignatureDef signature_def;
    for (const string& input : inputs) {
      (*signature_def.mutable_inputs())[input].set_name(input);

  /// SavedModel.
  virtual Session* GetSession() const = 0;

  /// Returns a map from signature name to SignatureDef for all signatures in
  /// in the SavedModel.
  virtual const protobuf::Map<string, SignatureDef>& GetSignatures() const = 0;
};

/// SavedModel representation once the SavedModel is loaded from storage.
///

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/framework/attr_value.pb.h"
#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;

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

/// Classification API constants.

/// Classification inputs.
static constexpr char kClassifyInputs[] = "inputs";

/// Classification method name used in a SignatureDef.
static constexpr char kClassifyMethodName[] = "tensorflow/serving/classify";

/// Classification classes output.
static constexpr char kClassifyOutputClasses[] = "classes";

/// Classification scores output.

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`.

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


def get_signatures_from_saved_model(
    saved_model_path: str,
    signature_keys: Optional[Sequence[str]] = None,
    tags: Optional[Collection[str]] = None,
) -> Dict[str, meta_graph_pb2.SignatureDef]:
  """Gets a map from signature keys to their SignatureDef.

  Args: