#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/debugger.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h"

  //     py_function_lib.py:save_exported_model,
  // )

  // Runs calibration on a model saved at `saved_model_path`. `exported_model`
  // should be the corresponding exported model resulting from the
  // pre-calibration step. `signature_keys` is a set of keys that identify a
  // SignatureDef to run the calibration on. `tags` is a set of strings that

) -> tuple[float, float]:
  """Calculates min and max from statistics using calibration options.

  Args:
    statistics: Collected calibration statistics.
    calib_opts: Calibration options used for calculating min and max.

  Returns:
    (min_value, max_value): Min and max calculated using calib_opts.

  Raises:
    ValueError: Unsupported calibration method is given.
  """
  calib_method = calib_opts.calibration_method

    statistics.histogram_statistics.hist_freq.extend(hist_freq.tolist())

    # Histogram calibration methods should remove outliers.
    min_value, max_value = calibration_algorithm.get_min_max_value(
        statistics, calibration_options
    )

    # Since the min/max values may differ slightly for each calibration method,
    # also check the nearby values.
    self.assertAllInRange(min_value, 49, 51)

    ],
) -> bool:
  """Runs calibration and adds calibration statistics to exported model.

  Args:
    saved_model_path: Path to the SavedModel to run calibration.
    signature_keys: List of signature keys corresponding to SignatureDefs to run
      calibration on.
    tags: A set of tags that identify the MetaGraphDef.
    force_graph_mode_calibration: If True, runs the calibration in graph mode.

using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorAverageMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorBase;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorHistogram;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsMap;

}  // namespace

namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::tensorflow::quantization::PyFunctionLibrary;
using CalibrationStatisticsFlatMap =
    absl::flat_hash_map<std::string, CalibrationStatistics>;

}  // namespace

// Reads the calibration statistics from the given directory.

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"

namespace tensorflow {
namespace calibrator {

// Abstract base class for CalibrationStatisticsCollcetor such as
// CalibrationStatisticsCollectorMinMax. Each class collects different
// statistics based on the calibration methods.
class CalibrationStatisticsCollectorBase {
 public:
  // Collect data for calibration.
  virtual void Collect(float min, float max,

namespace mlir::quant::stablehlo {

// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.
//

          clEnumValN(TEST_CASE_HISTOGRAM_PERCENTILE, "HISTOGRAM_PERCENTILE",
                     "Uses HISTOGRAM_PERCENTILE calibration method"),
          clEnumValN(TEST_CASE_HISTOGRAM_MSE_BRUTEFORCE,
                     "HISTOGRAM_MSE_BRUTEFORCE",
                     "Uses HISTOGRAM_MSE_BRUTEFORCE calibration method"),
          clEnumValN(TEST_CASE_HISTOGRAM_MSE_MAX_FREQUENCY,