}

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

namespace mlir {
namespace TF {

// Create a pass that groups ops into functions that only contain one dialect.
std::unique_ptr<Pass> CreateGroupByDialectPass();

// Register this pass in the global registry of MLIR.
void RegisterGroupByDialectPass();

}  // namespace TF
}  // namespace mlir

#include "tensorflow/core/tfrt/runtime/runtime.h"

namespace tensorflow {

class BackendCompiler {
 public:
  virtual ~BackendCompiler();

  virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}

  // Compile the `module` in TF dialect. The result module should be also in TF
  // dialect.
  virtual absl::Status CompileTensorflow(
      tfrt_stub::ModelRuntimeContext& model_context,

  tensorflow::InitMlir y(&argc, &argv);

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

  registry.addExtension(
      +[](mlir::MLIRContext *ctx, mlir::TF::TensorFlowDialect *dialect) {
        dialect->addInterfaces<TFReductionPatternInterface>();
      });

  mlir::MLIRContext context(registry);

  return failed(mlirReduceMain(argc, argv, context));

    : public impl::LowerVariableOpsToMlProgramPassBase<
          LowerVariableOpsToMlProgramPass> {
  explicit LowerVariableOpsToMlProgramPass() = default;
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<mlir::tf_saved_model::TensorFlowSavedModelDialect,
                    ml_program::MLProgramDialect>();
  }
  void runOnOperation() override {
    auto module = getOperation();

                                   FunctionLibraryDefinition* flib_def,
                                   Options options = Options());

  // Like `MarkForCompilation` but creates `flib_def` from the op registry.
  static Status MarkForCompilation(std::unique_ptr<Graph>* graph,
                                   Options options = Options());
};
}  // namespace tensorflow

  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

// 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.