llvm::SMLoc());
    auto module = mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, &ctx);
    if (!module) {
      return false;
    }

    mlir::SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &ctx);
    if (failed(mlir::verify(*module))) {
      module->emitError("Invalid MLIR module: failed verification.");
      return false;
    }
    return true;
  });

    context_.loadAllAvailableDialects();
    return absl::OkStatus();
  }

  Status LegalizeSingleOp(Operation& op) {
    SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);

    OpBuilder op_builder(&op);
    EmptyPatternRewriter pattern_rewriter(op_builder);

    LogicalResult result =
        Tf2XlaRewriter::RewriteOp(&op, pattern_rewriter,

  mlir::DialectRegistry registry;
  mlir::func::registerAllExtensions(registry);
  MLIRContext context(registry);
  llvm::SourceMgr source_mgr;
  mlir::SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &context);

  if (input_mlir) {
    // TODO(@zichuanwei): hack to enable mlir conversion via this tool, will get
    // back to do it properly in the future
    mlir::DialectRegistry registry;

  mlir::quant::RegisterOps();

  MLIRContext context(registry);
  context.loadAllAvailableDialects();

  llvm::SourceMgr source_mgr;
  SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &context);

  auto verbose_dir = llvm::sys::path::parent_path(output_path).str();

  const char* type_set[] = {"tflite", "mlir"};
  bool found = std::find(std::begin(type_set), std::end(type_set),