namespace py = ::pybind11;

namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::MinMaxValue;
using ::tensorflow::SignatureDef;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RepresentativeDatasetFile;

def clear_data_from_calibrator(id: bytes) -> None: ...

# LINT.ThenChange()

# LINT.IfChange(get_statistics_from_calibrator)
def get_statistics_from_calibrator(
    id: bytes,
) -> calibration_statistics_pb2.CalibrationStatistics: ...

                       absl::Span<const int64_t> histogram) = 0;

  virtual void ClearData() = 0;
  // Return the statistics needed for a given calibration method.
  virtual std::optional<CalibrationStatistics> GetStatistics() const = 0;
  virtual ~CalibrationStatisticsCollectorBase() = default;
};

}  // namespace calibrator
}  // namespace tensorflow

namespace tensorflow {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using CalibrationMethod =
    ::stablehlo::quantization::CalibrationOptions_CalibrationMethod;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorAverageMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorBase;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorHistogram;

    : public internal::SerializedProtobufCaster<tensorflow::GraphDef> {};

template <>
struct type_caster<tensorflow::calibrator::CalibrationStatistics>
    : public internal::SerializedProtobufCaster<
          tensorflow::calibrator::CalibrationStatistics> {};

template <>
struct type_caster<stablehlo::quantization::QuantizationConfig>
    : public internal::SerializedProtobufCaster<

  void ClearData() override;

  void Collect(float min, float max,
               absl::Span<const int64_t> histogram) override;

  std::optional<CalibrationStatistics> GetStatistics() const override;

 private:
  // Expands the histogram so the lower_bound and upper_bound can fit in the
  // histogram. Returns the indexes associated to those values.

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

// Adds calibrated min / max values to CustomAggregator nodes in `graph_def`.
// The min and max values will be added to the "min" and "max" attributes,

using ::tensorflow::AssetFileDef;
using ::tensorflow::SignatureDef;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;
using CalibrationStatisticsFlatMap =
    absl::flat_hash_map<std::string, CalibrationStatistics>;

}  // namespace

absl::Status RunCalibrationPasses(

    return cls

  return decorator


class _CalibrationAlgorithmBase(abc.ABC):
  """Abstract base class for calibration algorithm."""

  def __init__(
      self,
      statistics: calib_stats_pb2.CalibrationStatistics,
      calib_opts: stablehlo_quant_config_pb2.CalibrationOptions,
  ):
    self._statistics = statistics
    self._calib_opts = calib_opts

  @abc.abstractmethod

  // hinting and definition should also change.
  // LINT.IfChange(get_calibration_min_max_value)
  virtual std::optional<stablehlo::quantization::MinMaxValue>
  GetCalibrationMinMaxValue(const tensorflow::calibrator::CalibrationStatistics&
                                calibration_statistics,
                            const ::stablehlo::quantization::CalibrationOptions&
                                calibration_options) const = 0;