ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(
      kName, /*add_passes_func=*/
      [&config](PassManager& pm) {
        AddPreCalibrationPasses(pm, config.calibration_options(),
                                config.specs(), config.debugger_config());
      },
      *ctx_, module_op));
  return module_op;
}

// Adds passes for static-range quantization pre-calibration. Inserts ops
// required to collect tensor statistics.
void AddPreCalibrationPasses(
    OpPassManager& pm,
    const ::stablehlo::quantization::CalibrationOptions& calibration_options,
    const ::stablehlo::quantization::QuantizationSpecs& specs,
    const ::stablehlo::quantization::DebuggerConfig& debugger_config);

  // activation of static range quantization (SRQ). Quantization calibration
  // method is set to MIN_MAX by default.
  stablehlo.quantization.CalibrationOptions calibration_options = 15;

  // Configuration related to quantization debugger.
  stablehlo.quantization.DebuggerConfig debugger_config = 16;

  reserved 3;

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

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

#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

#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> {};

from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test

_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)


class CustomAggregatorTest(test.TestCase):

  def setUp(self):
    super(CustomAggregatorTest, self).setUp()
    ops.disable_eager_execution()

  def testBypassAndMinMax(self):