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>(CreateVerifyClusteringPass());
  }

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

 private:
  mlir::MLIRContext context_;

          return registry;
        }()) {
    context_.loadAllAvailableDialects();

    mlir::OpBuilder builder(&context_);
    module_ = builder.create<mlir::ModuleOp>(builder.getUnknownLoc());

    builder.setInsertionPointToStart(module_->getBody());
    auto func = builder.create<mlir::func::FuncOp>(  //
        builder.getUnknownLoc(), "main", builder.getFunctionType({}, {}));

  }

  mlir::LogicalResult Run() {
    mlir::OwningOpRef<mlir::ModuleOp> module_before = module_->clone();
    LogicalResult run_result = pm_->run(module_.get());

    EXPECT_TRUE(ModulesEqual(*module_before, *module_));
    return run_result;
  }

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

  explicit TacPass(const TacModule* module)
      : OperationPass<T>::OperationPass(mlir::TypeID::get<T>()),
        module_(module) {}

  ~TacPass() override = default;

  const TargetHardware* GetTargetHardware(
      const std::string& hardware_name) const {
    return module_ != nullptr
               ? module_->GetTargetHardware(hardware_name)
               : mlir::TFL::tac::GetTargetHardware(hardware_name);
  }

  }

 protected:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  llvm::SourceMgr source_manager_;
};

TEST_F(Tf2XlaRewriterTest, LegalizesOpWithTf2xlaHloImporter) {
  TF_EXPECT_OK(LegalizeModule());

  int num_tuple_ops = 0;
  module_->walk([&num_tuple_ops](TupleOp tuple_op) { num_tuple_ops += 1; });

  EXPECT_EQ(num_tuple_ops, 0);
}

 public:
  XlaLegalizeTargetsTest()
      : context_(GetDefaultDialectRegistry()),
        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(&context_))),
        builder_(&module_->getBodyRegion()) {
    context_.loadAllAvailableDialects();
  }

 protected:
  mlir::MLIRContext context_;
  mlir::OwningOpRef<mlir::ModuleOp> module_;
  mlir::OpBuilder builder_;
};

TEST_F(XlaLegalizeTargetsTest, CreatesConversionTargets) {

    TF_ASSERT_OK_AND_ASSIGN(module_,
                            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_;
};

  TargetAnnotationPass(const TargetAnnotationPass& copy)
      : TacFunctionPass(copy.module_) {}
  explicit TargetAnnotationPass(llvm::ArrayRef<std::string> device_specs)
      : TacFunctionPass(nullptr) {
    device_specs_flag_ = device_specs;
  }

  explicit TargetAnnotationPass(const TacModule* module)
      : TacFunctionPass(module) {}

 private:
  void runOnFunction() override;
  void SetTargetAnnotation(Operation* op,

  }

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

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

TEST_F(VerifyNoOutsideCompilationMarkersPassTest, PassesValidOps) {
  static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

        tensorflow::DeserializeMlirModule(module_string, &context_, &module_));

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