using ::mlir::quant::stablehlo::WeightOnlyPtqComponent;
using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::PopulateDefaults;
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::PyFunctionLibrary;

// Returns signature key -> `SignatureDef` mapping, excluding the signature for

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

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;

StaticRangePtqComponent::StaticRangePtqComponent(

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::quantization::RunPasses;

PreCalibrationComponent::PreCalibrationComponent(
    absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> PreCalibrationComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(

namespace {

using ::stablehlo::quantization::ExpandPresets;
using ::stablehlo::quantization::PopulateDefaults;
using ::stablehlo::quantization::QuantizationConfig;

class TestPreCalibrationComponentPass
    : public impl::TestPreCalibrationComponentPassBase<
          TestPreCalibrationComponentPass> {
 public:

using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::QuantizationSpecs;
using ::tensorflow::quantization::RunPasses;

PostCalibrationComponent::PostCalibrationComponent(
    absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> PostCalibrationComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(

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

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

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

// Runs weight-only quantization on a SavedModel at
// `src_saved_model_path` and saves the resulting model to
// `dst_saved_model_path`.

namespace mlir::quant::stablehlo {
namespace {

using ::stablehlo::quantization::ExpandPresets;
using ::stablehlo::quantization::PopulateDefaults;
using ::stablehlo::quantization::QuantizationConfig;
using ::testing::Contains;
using ::testing::SizeIs;
using ::testing::StartsWith;
using ::testing::StrEq;
using ::tsl::testing::IsOk;

// Matches an operation whose `getSymName` equals `name`.

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

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