const tensorflow::quantization::PyFunctionLibrary*
                      quantization_py_function_library = nullptr);

// Quantize the model with calibration data. Throw errors if `fully_quantize`
// is specified by the calibration data are not sufficient to quantize the
// model.
PyObject* MlirQuantizeModel(PyObject* data, bool disable_per_channel,
                            bool fully_quantize, int inference_type,

    targets.
  }];
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}

def InsertCalibrationStatisticsSaverPass : Pass<"stablehlo-insert-calibration-statistics-saver", "ModuleOp"> {
  let summary = "Inserts `CalibrationStatisticsSaver` op to collect and save calibration statistics.";
  let description = [{
    Finds all `CustomAggregator` ops in the each function and add a single

  const int32_t bin_index = CalculateBinIndex(value, lower_bound, bin_width);
  return std::clamp(bin_index, 0, num_bins - 1);
}

// Checks if the given method is a histogram-based calibration method.
inline bool IsHistogramCalibration(
    const CalibrationOptions::CalibrationMethod method) {
  return method ==
             CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE ||
         method ==

syntax = "proto3";

package tensorflow.calibrator;

option cc_enable_arenas = true;

// Calibration algorithm's collecting statistics.
// NEXT_ID: 4
message CalibrationStatistics {
  message MinMaxStatistics {
    // global minimum of all sample datasets.
    float global_min = 1;

    // global maximum of all sample datasets.
    float global_max = 2;
  }

  message AverageMinMaxStatistics {

  # Set `force_graph_mode_calibration` to True to avoid skipping op execution,
  # which are not connected to return ops, during calibration execution.
  # TODO: b/335031954 - Bring back support to run calibration in Eager mode.
  logging.debug(
      'Setting `force_graph_mode_calibration = True` to ensure the calibration'
      ' mode is executed properly.'
  )
  quantization_options.force_graph_mode_calibration = True

#define EIGEN_USE_THREADS
#include <cstdint>
#include <string>

#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

// RUN: stablehlo-quant-opt %s -stablehlo-test-post-calibration-component \
// RUN:   -split-input-file | FileCheck %s
// RUN: stablehlo-quant-opt %s -stablehlo-test-post-calibration-component='unpack-quantized-types=false' \
// RUN:   -split-input-file | FileCheck %s --check-prefix=CHECK-NO-UNPACK

// Tests that a simple dot_general (lifted as a function) with CustomAggregators
// around it is quantized. The resulting graph has quantized types unpacked into
// int ops.

absl::StatusOr<mlir::quant::QuantizedType> GetQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder,
    bool is_constant = false, mlir::Type storage_type = {});

// Imports float tensor with calibration value into calibrated quantized type.
absl::StatusOr<mlir::quant::QuantizedType> GetCalibratedQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder);

#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"

    TensorOf<[TF_Float32]>:$max,
    TensorOf<[TF_Int64]>:$histogram
  );
}

def TF_CalibrationStatisticsSaverOp : TF_Op<"CalibrationStatisticsSaver", []> {
  let summary = "Aggregates and saves calibration statistics.";

  let arguments = (ins
    Variadic<TensorOf<[TF_Float32, TF_Int64]>>:$inputs,

    StrAttr:$output_file_path,
    StrArrayAttr:$ids,
    I32ArrayAttr:$calibration_methods
  );