// 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);

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

  GetCalibrationMinMaxValue(const tensorflow::calibrator::CalibrationStatistics&
                                calibration_statistics,
                            const ::stablehlo::quantization::CalibrationOptions&
                                calibration_options) const = 0;
  // LINT.ThenChange(
  //     pywrap_function_lib.pyi:get_calibration_min_max_value,
  //     py_function_lib.py:get_calibration_min_max_value,

  // Path to save debugger related logs. Defaults to '/tmp/dumps'.
  string log_dir_path = 3;
}

// Defines various calibration options.
// Next ID: 6
message CalibrationOptions {
  // Configurations for calibration methods.
  // Next ID: 7
  enum CalibrationMethod {
    CALIBRATION_METHOD_UNSPECIFIED = 0;
    // Use the min, max values of all sample datasets.

def _populate_calibration_options(
    quantization_options: quant_opts_pb2.QuantizationOptions,
):
  """Populates default values for CalibrationOptions.

  Args:
    quantization_options: An instance of QuantizationOptions with a field
      specifying CalibrationOptions
  """

  calib_opts = quantization_options.calibration_options
  if (
      calib_opts.calibration_method

// Inserts custom aggregation operators for the calibration procedure.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateInsertCustomAggregationOpsPass(
    const ::stablehlo::quantization::CalibrationOptions& calib_opts);

// Replaces composite functions with quantized composite functions. After this
// pass runs, functions in the given graph will be replaced with their quantized

    Args:
      calibration_statistics_serialized: Serialized `CalibrationStatistics`.
        This will be the source to calculate min and max values from.
      calibration_options_serialized: Serialized `CalibrationOptions`. Specifies
        how the min / max should be calculated.

    Returns:
      (min_value, max_value): Min and max calculated using calib_opts. `None`
      upon error.
    """

  // Defines calibration options for quantization. This option is only used for
  // 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;