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

    ],
) -> 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.

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,

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

namespace stablehlo::quantization {

// Reads the calibration statistics from the given directory.
absl::StatusOr<absl::flat_hash_map<
    std::string, tensorflow::calibrator::CalibrationStatistics>>
ReadStatistics(absl::string_view calibration_data_dir);

  string unquantized_dump_model_path = 2;

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

                  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