from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2 as calib_stats_pb2


_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)
_REGISTRY = {}


def _implements(calib_method: _CalibrationMethod):
  def decorator(cls):
    assert calib_method not in _REGISTRY
    _REGISTRY[calib_method] = cls
    return cls

#include "tsl/platform/path.h"
#include "tsl/platform/statusor.h"

namespace stablehlo::quantization {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::tensorflow::quantization::PyFunctionLibrary;
using CalibrationStatisticsFlatMap =

#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;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {
namespace calibrator {

using ::stablehlo::quantization::CalibrationOptions;

// AverageMinMax calibration calculates the average of min and max values.
// average of min = sum of min values / number of samples
// average of max = sum of max values / number of samples

    %1 = "tf.Relu6"(%0) : (tensor<*xf32>) -> tensor<*xf32>
    func.return %1 : tensor<*xf32>
  }
}

// CalibrationOptions(calibration_method=CALIBRATION_METHOD_MIN_MAX)
// MIN-MAX-CHECK: func @wrap_composite_func

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {
namespace calibrator {

using ::stablehlo::quantization::CalibrationOptions;

// MinMax calibration calculates the global min and global max values.
// global min = min of given sample inputs
// global max = max of given sample inputs
class CalibrationStatisticsCollectorMinMax

          tensorflow::quantization::QuantizationOptions> {};

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

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

              stablehlo_quant_config_pb2.CalibrationOptions(
                  calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX
              ),
              stablehlo_quant_config_pb2.CalibrationOptions(
                  calibration_method=_CalibrationMethod.CALIBRATION_METHOD_AVERAGE_MIN_MAX
              ),
              stablehlo_quant_config_pb2.CalibrationOptions(

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::DebuggerConfig;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationSpecs;

void AddPreCalibrationPasses(OpPassManager& pm,
                             const CalibrationOptions& calibration_options,
                             const QuantizationSpecs& quantization_specs,

// max values.
absl::Status AddCalibrationStatistics(
    mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
    const stablehlo::quantization::CalibrationOptions& calibration_options,
    const tensorflow::quantization::PyFunctionLibrary& py_function_library);

// Checks if the model required calibration.
bool IsCalibrationRequired(mlir::ModuleOp module_op);