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

      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;

  std::set<std::string> keys;
  for (const auto& it : map) {
    keys.insert(it.first);
  }
  return keys;
}

Status GetInputValues(
    const SignatureDef& signature,
    const ::google::protobuf::Map<std::string, ::tensorflow::TensorProto>& request_inputs,
    std::vector<std::pair<string, Tensor>>& inputs) {
  const TensorProto* tensor_proto;

    const std::vector<std::string>& signature_keys,
    const std::unordered_set<std::string>& tags,
    const QuantizationOptions& quantization_options,
    const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
    const PyFunctionLibrary& py_function_library,
    const absl::flat_hash_map<std::string, RepresentativeDatasetFile>&
        representative_dataset_file_map_serialized);

#include "tensorflow/core/platform/protobuf_internal.h"

namespace tensorflow {
namespace internal {

// A SavedModel may store the name of the initialization op to run in the
// in the SignatureDef (v2) or a collection (v1). If an init_op collection
// exists, then the collection must contain exactly one op.
Status GetInitOp(const string& export_dir, const MetaGraphDef& meta_graph_def,

// Loads a SavedModel at `saved_model_path` and converts it to `mlir::ModuleOp`.
//
// `tags` identify the `tensorflow::MetaGraphDef` to load from the SavedModel.
// Similarly, `signature_keys` identify the functions (`SignatureDef`s) to load
// within the `MetaGraphDef`. `ctx` is the `MLIRContext`, which should outlive
// the returned `ModuleOp`, thus marked with the lifetime bound attribute.

    const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info);

// Hashes the contents of `signature_def`.
absl::StatusOr<uint64_t> HashSignatureDef(
    const Map<std::string, ::tensorflow::SignatureDef>& signature_def_map,
    const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
    riegeli::RecordReader<riegeli::FdReader<>>& reader,
    const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info);

  // (TF1 SavedModel only) Collection of tags identifying the MetaGraphDef
  // within the SavedModel to analyze. If not specified, ["serve"] is used.
  repeated string tags = 5;

  // Sequence of keys identifying SignatureDef containing inputs and outputs.
  // If not specified, ["serving_default"] is used.
  repeated string signature_keys = 6;

  // A map from signature keys to the corresponding representative dataset.