}

bool BuildAndRunPipeline(absl::string_view module_string,
                         const std::function<void(PassManager*)>& passes) {
  mlir::registerPassManagerCLOptions();
  MLIRContext context;

  OwningOpRef<ModuleOp> module =
      GetMlirModuleFromString(module_string, &context).value();

  PassManager pm(&context);

  if (mlir::failed(mlir::applyPassManagerCLOptions(pm))) return false;
  passes(&pm);

#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "mlir/Pass/PassManager.h"  // from @llvm-project

namespace tensorflow {
namespace quantization {

// Enables IR printing for `pm`. When the passes are run, each pass will dump to
// its own file with prefix `file_name_prefix`.
void EnableIrPrinting(mlir::PassManager &pm,
                      absl::string_view file_name_prefix);

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "mlir/Support/TypeID.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"

#include <optional>

#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tsl/platform/errors.h"

#include <memory>
#include <string>

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

using mlir::PassManager;

// Setup the input pass manager to enable IR dumping after each pass.

#include <string>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/save_variables.h"

    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_;
};

TEST_F(VerifyClusteringPassTest, OnlyTfFunctionalPasses) {

#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"

#include <memory>
#include <optional>

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tsl/platform/errors.h"

#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"

namespace stablehlo {
namespace quantization {
// Adds passes for quantization of individual quantizable components.
// (i.e. activation, weight, bias)
void AddQuantizationPasses(mlir::PassManager& pass_manager,