registry.insert<mlir::func::FuncDialect>();
          registry.insert<mlir::TFL::TensorFlowLiteDialect>();
          return registry;
        }()) {
    context_.loadAllAvailableDialects();

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

    builder.setInsertionPointToStart(module_->getBody());

    dialects.insert<mlir::BuiltinDialect, mlir::func::FuncDialect,
                    mlir::stablehlo::StablehloDialect>();
    ctx_ = std::make_unique<mlir::MLIRContext>(dialects);
    ctx_->loadAllAvailableDialects();
  }

  void TearDown() override {
    // Delete files in the test directory.
    std::vector<std::string> files;
    TF_ASSERT_OK(

                  mlir::TFR::TFRDialect>();
  // clang-format on
  mlir::func::registerAllExtensions(registry);
  mlir_ctx->appendDialectRegistry(registry);
  mlir_ctx->loadAllAvailableDialects();

  // Load the TFR functions in a mlir::ModuleOp
  auto memory_buffer = llvm::MemoryBuffer::getMemBuffer(
      llvm::StringRef(tfr_raw_text.data(), tfr_raw_text.size()));
  llvm::SourceMgr source_mgr;