using ::tensorflow::quantization::RunPasses;
using CalibrationStatisticsFlatMap =
    absl::flat_hash_map<std::string, CalibrationStatistics>;

}  // namespace

absl::Status RunCalibrationPasses(
    mlir::ModuleOp module_op, MLIRContext& ctx,
    absl::string_view calibration_data_dir,
    const bool force_regenerate_calibration_data) {
  // Disable DumpTensor ops when running calibration.

  // Signature keys to identify the functions to load & quantize.
  const std::vector<std::string> signature_keys_;
};

// Runs passes to prepare the calibration model.
absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,
                                  absl::string_view calibration_data_dir,
                                  bool force_regenerate_calibration_data);

using ::mlir::quant::stablehlo::kExportStepSuffix;
using ::mlir::quant::stablehlo::PostCalibrationComponent;
using ::mlir::quant::stablehlo::PreCalibrationComponent;
using ::mlir::quant::stablehlo::RunCalibrationPasses;
using ::mlir::quant::stablehlo::UpdateFunctionAliases;
using ::mlir::quant::stablehlo::WeightOnlyPtqComponent;
using ::stablehlo::quantization::AddCalibrationStatistics;