syntax = "proto3";

package tensorflow.quantization;

import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/protobuf/meta_graph.proto";
import "tensorflow/core/protobuf/saver.proto";

option cc_enable_arenas = true;

// Represents an exported TensorFlow model. It consists of a GraphDef and extra
// metadata required for building a SavedModel. This message is primarily used

  // Retrieve a SignatureDefFunction from a SavedModel, using the key of the
  // SignatureDef map:
  // https://github.com/tensorflow/tensorflow/blob/69b08900b1e991d84bce31f3b404f5ed768f339f/tensorflow/core/protobuf/meta_graph.proto#L89
  virtual Status GetSignatureDefFunction(const std::string& signature_def_key,
                                         SignatureDefFunction** function) = 0;

  virtual SavedModelV2Bundle* GetBundle() = 0;

#include <utility>
#include <vector>

#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"

namespace tensorflow {
namespace saved_model {

// Utility functions for SavedModel reading and writing.

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace stablehlo::quantization::pywrap {

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

#include "absl/status/statusor.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"

namespace tensorflow {
Status ReadSavedModel(absl::string_view export_dir,
                      SavedModel* saved_model_proto);

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace stablehlo::quantization::pywrap {

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

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {

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

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace mlir::quant::stablehlo {

// Performs int8 weight-only quantization on dot_general ops.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace py = ::pybind11;

namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::MinMaxValue;
using ::tensorflow::SignatureDef;

// In both cases, a SignatureDef is serialized as a SignatureDef protobuf:
// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/core/protobuf/meta_graph.proto#L260-L330
// and represents a computation defined by a TF subgraph.
// These Signatures were primarily designed to be interoperable with the legacy
// TF 1 Session-based C++ SavedModelBundle loading APIs: