CalibrationOptions::CALIBRATION_METHOD_MIN_MAX);
        break;
      case TEST_CASE_AVERAGE_MIN_MAX:
        calib_opts_.set_calibration_method(
            CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX);
        break;
      case TEST_CASE_HISTOGRAM_PERCENTILE: {
        calib_opts_.set_calibration_method(
            CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE);

  EXPECT_EQ(statistics.value().average_min_max_statistics().num_samples(), 1);
}

TEST(HistogramStatisticsCollectorTest, SingleBatchSimple) {
  CalibrationOptions calib_opts;
  calib_opts.set_calibration_method(
      CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
  auto collector = CalibrationStatisticsCollectorHistogram();

  collector.Collect(

    switch (calibration_method) {
      case CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX:
        collector =
            std::make_unique<CalibrationStatisticsCollectorAverageMinMax>();
        break;
      case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE:
      case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE:
      case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC:

inline bool IsHistogramCalibration(
    const CalibrationOptions::CalibrationMethod method) {
  return method ==
             CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE ||
         method ==
             CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE ||
         method == CalibrationOptions::
                       CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY ||

}

TEST_F(CalibrationStatisticsSaverTest, MultipleStats) {
  std::vector<std::string> ids{"1", "2"};
  std::vector<int32_t> calibration_methods{
      CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX,
      CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE};

  std::vector<NodeDefBuilder::NodeOut> inputs;
  inputs.emplace_back("min", 0, DT_FLOAT);
  inputs.emplace_back("max", 0, DT_FLOAT);

namespace {

using ::stablehlo::quantization::CalculateBinIndexSafe;
using ::stablehlo::quantization::CalculateBinWidth;
using ::stablehlo::quantization::CalculateLowerBound;
using ::stablehlo::quantization::CalibrationOptions;
using CPUDevice = ::Eigen::ThreadPoolDevice;
using CalibrationMethod =
    ::stablehlo::quantization::CalibrationOptions_CalibrationMethod;

}  // namespace

REGISTER_OP("CustomAggregator")

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