if (include_extensions) {
    mlir::func::registerAllExtensions(registry);
  }
}

// Inserts all the TensorFlow dialects in the provided registry. This is
// intended for tools that need to register dialects before parsing .mlir files.
inline void RegisterAllTensorFlowDialects(DialectRegistry &registry) {
  RegisterAllTensorFlowDialectsImpl(registry, true);
}
}  // namespace mlir

      "tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir");

  mlir::DialectRegistry registry;
  RegisterTFRTDialects(registry);
  registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
  registry.insert<tfrt::fallback_sync::FallbackSyncDialect>();

  mlir::MLIRContext context(registry);
  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

    MlirToFlatBufferFileTranslateFunction, [](DialectRegistry& registry) {
      registry.insert<quant::QuantizationDialect,
                      quantfork::QuantizationForkDialect>();
      mlir::RegisterAllTensorFlowDialects(registry);
      registry.insert<TFL::TensorFlowLiteDialect>();
      registry.insert<arith::ArithDialect>();
      registry.insert<func::FuncDialect>();
      registry.insert<mlir::vhlo::VhloDialect>();

PYBIND11_MODULE(tfr_wrapper, m) {
  m.def("verify", [](std::string input) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::arith::ArithDialect, mlir::scf::SCFDialect,
                    mlir::TF::TensorFlowDialect, mlir::func::FuncDialect,
                    mlir::shape::ShapeDialect, mlir::TFR::TFRDialect>();
    mlir::MLIRContext ctx(registry);
    ctx.loadAllAvailableDialects();

    llvm::SourceMgr source_mgr = llvm::SourceMgr();

  mlir::DialectRegistry registry;
  mlir::registerAllDialects(registry);
  mlir::RegisterAllTensorFlowDialects(registry);

  mlir::registerAllPasses();
  mlir::registerTensorFlowPasses();

  EXPECT_TRUE(mlir::MlirOptMain(output_stream->os(), std::move(input_file),
                                registry,
                                mlir::MlirOptMainConfig{}

  for (int i = 0; i < grad_inputs.size(); i++) {
    grad_inputs[i] = nullptr;
  }
  return absl::OkStatus();
}

Status RegisterNotDifferentiable(GradientRegistry* registry, const string& op) {
  return registry->Register(op, [](const ForwardOperation& op) {
    return new NotDifferentiableGradientFunction;
  });
}
}  // namespace gradients

  tensorflow::tfrt_compiler::
      RegisterNonTPULowerClusterToRuntimeOpsPassPipeline();

  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);

  return failed(
      mlir::MlirOptMain(argc, argv, "TensorFlow pass driver\n", registry));

}

static TranslateFromMLIRRegistration mlir_to_graph_translate(
    /*name=*/"mlir-to-graph", /*description=*/"convert mlir to graph",
    MlirToGraphTranslateFunction, [](DialectRegistry& registry) {
      mlir::RegisterAllTensorFlowDialects(registry);
    });

static LogicalResult MlirToGraphdefTranslateFunction(
    ModuleOp module, llvm::raw_ostream& output) {
  if (!module) return failure();

  func.return %3 : tensor<2x1xf32>
})";
  const std::string kExpectedFB = CreateRuntimeMetadata();
  mlir::DialectRegistry registry;
  registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
                  mlir::func::FuncDialect>();
  mlir::MLIRContext context(registry);
  auto module = mlir::OwningOpRef<mlir::ModuleOp>(
      mlir::parseSourceString<mlir::ModuleOp>(kMLIR, &context));

namespace mlir {

void RegisterCommonToolingDialects(mlir::DialectRegistry& registry) {
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::mhlo::registerAllMhloDialects(registry);
  mlir::registerAllDialects(registry);
  mlir::registerAllExtensions(registry);
  mlir::stablehlo::registerAllDialects(registry);

  registry.insert<mlir::TFL::TensorFlowLiteDialect>();
  registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();