if (parser.parseKeyword(&keyword)) return Type();

  if (keyword == "tf_tensor") return TFTensorType::get(getContext());
  if (keyword == "tf_allocator") return TFAllocatorType::get(getContext());

  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
  return Type();
}

/// Print a type registered to this dialect.
void FallbackDialect::printType(Type type, DialectAsmPrinter &os) const {

    auto func =
        llvm::dyn_cast_or_null<func::FuncOp>(symbol_table.lookupSymbolIn(
            module, builder.getStringAttr(function_name)));
    if (func == nullptr) {
      return module.emitError()
             << "Cannot find function '" << function_name << "'";
    }

    const WalkResult result = WalkReachableFunctions(
        func,
        [&](func::FuncOp f) {

}

void ReplicaIDToDeviceOrdinalPass::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

  // Get the number of devices per host.
  int device_num = 0;
  mlir::TF::RuntimeDevices devices;

                                    BlockArgument& block_arg) {
  // check that inputs length is same as num_replicas.
  if (inputs.size() != replicate.getN()) {
    return replicate.emitError()
           << "Expected numper of inputs (" << inputs.size()
           << ") to append to replicate to be num_replicas ("
           << replicate.getN() << ")";
  }

    if (!sym) continue;
    auto callee = symtab.lookup<func::FuncOp>(sym.getRootReference());
    if (!callee) {
      // This is not expected to happen in practice.
      return op->emitError()
             << "Cannot find function " << sym.getRootReference();
    }
    callees.push_back(callee);
  }
  return success();
}

bool HasSingleBlock(func::FuncOp func) {

    if (op.intermediates.size() != 5) {
      auto err = errors::InvalidArgument(
          "operator has intermediate tensors but the number of them is not "
          "five.");
      return emitError(loc, err.ToString()), err;
    }
    // Create intermediate value

    const llvm::SmallVector<llvm::StringRef, 5> kIntermediateNames = {
        "input_to_input_intermediate", "input_to_forget_intermediate",

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MockSuccessPass)

  explicit MockSuccessPass() = default;

 private:
  void runOnOperation() override {
    getOperation().walk([](Operation* nestedOp) {
      nestedOp->emitError()
          << "Error at " << nestedOp->getName().getStringRef().str() << " op";
    });
  };
};

// MockFailurePass reports errors and fails.
class MockFailurePass

// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {
  auto launch_op = llvm::dyn_cast<tf_device::LaunchOp>(op->getParentOp());
  if (!launch_op) {
    op->emitError() << "TPUCopyWithDynamicShapeOp is not in a launch";
  }
  std::string device_str = launch_op.getDeviceAttr().getValue().str();
  std::string cpu0_device;
  if (failed(tensorflow::GetNonReplicatedCPU0(op, &cpu0_device)))

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

    if (!replication_info_attr)
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    auto replication_info_attr_str =
        mlir::dyn_cast<StringAttr>(replication_info_attr);
    if (!replication_info_attr_str ||
        replication_info_attr_str.getValue().empty())
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    // Remove `name` attribute.