// https://www.tensorflow.org/tutorials/load_data/tfrecord for details on the
// TFRecord format.
// Next ID: 2
message TfRecordFile {
  string path = 1;
}

// Configures a single representative dataset used to calibrate a single
// function.
// Next ID: 3
message RepresentativeDatasetConfig {
  oneof file {
    // Represents representative dataset saved as a .tfrecord file format.
    TfRecordFile tf_record = 1;
  }

from tensorflow.compiler.mlir.quantization.tensorflow import quantization_options_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_algorithm
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.python import pywrap_function_lib

#include "tsl/platform/statusor.h"

namespace stablehlo::quantization {
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

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

// CHECK-SAME: effective_hidden_scale_intermediate = tensor<*x!quant.calibrated<f32<-5.000000e-01:5.000000e-01>>>
// CHECK-SAME: input_to_cell_intermediate = tensor<*x!quant.calibrated<f32<-4.000000e+00:4.000000e+00>>>
// CHECK-SAME: input_to_forget_intermediate = tensor<*x!quant.calibrated<f32<-1.600000e+01:1.600000e+01>>>
// CHECK-SAME: input_to_input_intermediate = tensor<*x!quant.calibrated<f32<-3.200000e+01:3.200000e+01>>>

  // Runs calibration on `module_op` and returns a calibrated ModuleOp with
  // calibrated statistics embedded.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // Exports `module_op` to SavedModel at `dst_saved_model_path`. This is used
  // to export the pre-calibrated `module_op` to SavedModel so that the

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

      effective_hidden_scale_intermediate = tensor<!quant.calibrated<f32<-5.000000e-01:5.000000e-01>>>,
      fused_activation_function = "TANH",
      input_to_cell_intermediate = tensor<!quant.calibrated<f32<-4.000000e+00:4.000000e+00>>>,
      input_to_forget_intermediate = tensor<!quant.calibrated<f32<-1.600000e+01:1.600000e+01>>>,
      input_to_output_intermediate = tensor<!quant.calibrated<f32<-1.000000e+00:1.000000e+00>>>,

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

namespace tensorflow {
namespace calibrator {


class CalibrationStatisticsCollectorHistogram

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