GetMlirModuleFromString(module_string, &context_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addPass(CreateInferenceMetricsPass());
  }
  mlir::LogicalResult Run() { return pm_->run(module_.get()); }

 private:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

TEST_F(InferenceMetricsPassTest, RecordsTrueForTPUPartitionedCallOp) {

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(CreateVerifyClusteringPass());
  }

  mlir::LogicalResult Run() { return pm_->run(module_.get()); }

 private:
  mlir::MLIRContext context_;
  mlir::OwningOpRef<mlir::ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<func::FuncOp>(
        CreateVerifyNoOutsideCompilationMarkersPass());
  }

  mlir::LogicalResult Run() { return pm_->run(module_.get()); }

 private:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));
    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(CreateInputLoweringMetricsPass());
  }

  bool ModulesEqual(const ModuleOp& module_before,
                    const ModuleOp& module_after) {

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(
        quant::stablehlo::CreateConvertTFQuantTypesPass());
  }
  mlir::LogicalResult Run() { return pm_->run(module_.get()); }

 private:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(
        mlir::mhlo::CreateVerifyTFXLALegalizationPass(/*legalize_chlo=*/false));
  }
  mlir::LogicalResult Run() { return pm_->run(module_.get()); }

 private:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

  if (failed(pm_.run(user_module))) {
    return errors::Internal("Failed to run the decompose passes.");
  }
  return absl::OkStatus();
}

// Constructor of the decompose context.
TFRDecomposeContext::TFRDecomposeContext(mlir::ModuleOp tfr_module)
    : tfr_module_(tfr_module), pm_(tfr_module_.getContext()) {
  mlir::OpPassManager& func_pm = pm_.nest<mlir::func::FuncOp>();

  Status DecomposeGraph(mlir::ModuleOp user_module);

  // Erases the tfr_module created.
  void Destroy();

 private:
  mlir::ModuleOp tfr_module_;
  mlir::PassManager pm_;

  GraphExportConfig export_confs_;
};

// Decomposes the NodeDef to a set of primitive ops according to the decompose
// library loaded. Wrap the decomposed result in a FunctionDef.