}

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

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

  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));

}

void RuntimeFallbackExecutor::Prepare(llvm::StringRef mlir_input) {
  // We only support IR written in the Tensorflow dialect.
  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);

  llvm::SourceMgr source_mgr;
  source_mgr.AddNewSourceBuffer(
      llvm::MemoryBuffer::getMemBuffer(mlir_input, "test_ir"), llvm::SMLoc());

  mlir::DialectRegistry registry;
  registry.insert<mlir::func::FuncDialect, mlir::stablehlo::StablehloDialect,
                  mlir::TFL::TFLDialect, mlir::arith::ArithDialect,
                  mlir::TF::TensorFlowDialect, mlir::chlo::ChloDialect>();

  return failed(
      mlir::MlirOptMain(argc, argv, "ODML Converter Driver\n", registry));

  }

  if (input_mlir) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
                    mlir::func::FuncDialect>();
    mlir::func::registerAllExtensions(registry);
    context->appendDialectRegistry(registry);
    source_mgr->AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());
    return mlir::OwningOpRef<mlir::ModuleOp>(

    // This is a brief description of the pass.
    return "Convert tf.CustomAggregator op to quant.Stats";
  }

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<TF::TensorFlowDialect>();
    registry.insert<quant::QuantizationDialect>();
    registry.insert<quantfork::QuantizationForkDialect>();
  }

  void runOnOperation() override;
};

class ConvertCustomAggregationOpToQuantStats

                               std::vector<TensorShape>& arg_shapes,
                               tpu::TPUCompileMetadataProto& metadata_proto) {
  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));

  llvm::cl::ParseCommandLineOptions(argc, argv,
                                    "TF Saved Model to Stablehlo converter\n");

  mlir::DialectRegistry registry;
  RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  bool is_input_mlir_module = absl::EndsWith(input_path, ".mlir");
  std::vector<std::string> exported_model_signatures_in_vector =

// Sparsifies model to encode sparse tensors with proper format. Throws error if
// sparsification fails.
PyObject* MlirSparsifyModel(PyObject* data);

// Registers the given custom opdefs to TensorFlow global op registry.
PyObject* RegisterCustomOpdefs(PyObject* list);

// Returns the collected TFLite conversion errors.
std::vector<std::string> RetrieveCollectedErrors();

// Returns MLIR string dump of the given Flatbuffer model.