// Runs the static-range post-training quantization (PTQ) on `module_op`.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // A non-owning `MLIRContext`. This `MLIRContext` should exceed the lifetime
  // of `StaticRangePtqComponent`.
  absl::Nonnull<MLIRContext*> ctx_;

          tensorflow::calibrator::CalibrationStatistics> {};

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

template <>
struct type_caster<tensorflow::quantization::RepresentativeDatasetFile>
    : public internal::SerializedProtobufCaster<

  // should provide information necessary to configure the action's behavior.
  virtual absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) = 0;
};

}  // namespace mlir::quant::stablehlo

absl::StatusOr<mlir::ModuleOp> RunQuantization(
    const SavedModelBundle* saved_model_bundle,
    absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& saved_model_tags,
    const stablehlo::quantization::QuantizationConfig& quantization_config,
    const tensorflow::quantization::PyFunctionLibrary*
        quantization_py_function_lib,
    mlir::ModuleOp module_op);

}  // namespace tensorflow

#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::MLIRImportOptions;
using ::tensorflow::SavedModelBundle;
using ::tensorflow::SavedModelSignatureDefsToMlirImport;
using ::tensorflow::quantization::PreprocessAndFreezeGraph;

  explicit PreCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  absl::Nonnull<MLIRContext*> ctx_;
};

}  // namespace mlir::quant::stablehlo

using ::stablehlo::quantization::CreateRepresentativeDatasetFileMap;
using ::stablehlo::quantization::DisableDebugging;
using ::stablehlo::quantization::IsCalibrationRequired;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::ReadStatistics;
using ::stablehlo::quantization::RepresentativeDatasetConfig;
using ::stablehlo::quantization::io::CreateTmpDir;
using ::stablehlo::quantization::io::GetLocalTmpFileName;

  explicit PostCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

  void AddPasses(
      OpPassManager& pm,
      const ::stablehlo::quantization::QuantizationSpecs& specs,
      const ::stablehlo::quantization::PipelineConfig& pipeline_config) const;

absl::StatusOr<OwningOpRef<ModuleOp>> ImportSavedModel(
    absl::string_view saved_model_path,
    const std::vector<std::string>& signature_keys,
    const std::unordered_set<std::string>& tags,
    const ::stablehlo::quantization::QuantizationConfig& quantization_config,
    absl::string_view mlir_dump_file_prefix,
    absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
    MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND);

using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerInitType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::io::GetLocalTmpFileName;
using ::tensorflow::AssetFileDef;
using ::tensorflow::ConvertMlirToGraph;
using ::tensorflow::FunctionDefLibrary;
using ::tensorflow::FunctionLibraryDefinition;