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(
      kName, /*add_passes_func=*/

// lifted quantizable functions.
class PreCalibrationComponent : public Component {
 public:
  // Name of the post-training quantization pre-calibration step. Used for
  // debugging purposes.
  static constexpr absl::string_view kName = "quant_ptq_pre_calibration";

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

  absl::StatusOr<ModuleOp> Run(
      ModuleOp,

    absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases)
    : ctx_(ctx) {
  // Initialize the three sub-components.
  sub_components_[0] = std::make_unique<PreCalibrationComponent>(ctx_);
  sub_components_[1] = std::make_unique<CalibrationComponent>(
      ctx_, py_function_library, src_saved_model_path,
      std::move(function_aliases), std::move(tags),

};

void TestPreCalibrationComponentPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext& ctx = getContext();

  // Simply runs the PreCalibrationComponent with a default configuration.
  PreCalibrationComponent component(&ctx);
  QuantizationConfig quantization_config{};
  quantization_config.mutable_static_range_ptq_preset();
  quantization_config = ExpandPresets(PopulateDefaults(quantization_config));

  // Name of this component. Used for debugging purposes.
  static constexpr absl::string_view kName = "quant_static_range_ptq";

  // Constructs `StaticRangePtqComponent` by creating three sub-components:
  // `PreCalibrationComponent`, `CalibrationComponent`, and
  // `PostCalibrationComponent`. These are stored in `sub_components_` in
  // sequence. All arguments except `ctx` is used to initialize

include "mlir/Pass/PassBase.td"

def TestPreCalibrationComponentPass : Pass<"stablehlo-test-pre-calibration-component", "mlir::ModuleOp"> {
  let summary = "Test-only pass to test the PreCalibrationComponent.";
  let description = [{
    Runs the pre calibration passes for post-training quantization with default
    configuration.
  }];
  let dependentDialects = [

}

using PreCalibrationComponentTest = ::mlir::quant::QuantizationTestBase;

TEST_F(PreCalibrationComponentTest,
       HasCustomAggregatorOpAndQuantizableFuncForSimpleDotGeneral) {
  PreCalibrationComponent component(ctx_.get());
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {} {
      func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> attributes {} {

using ::mlir::quant::stablehlo::GetFunctionAliases;
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;