auto sym = mlir::dyn_cast<SymbolRefAttr>(attr.getValue());
    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 (!nested_module) return;

  InlinerInterface inliner(&getContext());
  auto walk_result = getOperation().walk([&](TF::PartitionedCallOp call_op) {
    if (!call_op.getF().getRootReference().getValue().starts_with(
            kNestedModule))
      return WalkResult::advance();
    // This is a call we need to inline!
    LLVM_DEBUG(llvm::dbgs()

    CallInterfaceCallable callable = call.getCallableForCallee();
    auto sym = callable.get<mlir::SymbolRefAttr>();
    EncapsulatePartitionedCall(pcall_op, sym.getRootReference());
  }
  // Partitioned calls are executed asynchronous. The calls outside of
  // device clusters therefore should not be inlined to perserve run-time
  // performance.

//===----------------------------------------------------------------------===//

LogicalResult ClusterFuncOp::verifySymbolUses(
    mlir::SymbolTableCollection& symbolTable) {
  StringAttr func_attr = getFuncAttr().getRootReference();
  func::FuncOp func =
      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, func_attr);
  if (!func) {
    return emitError("'func' attribute refers to an undefined function: ")

//===----------------------------------------------------------------------===//

LogicalResult BatchFunctionOp::verifySymbolUses(
    SymbolTableCollection& symbolTable) {
  StringAttr func_attr = getFAttr().getRootReference();
  func::FuncOp func =
      symbolTable.lookupNearestSymbolFrom<func::FuncOp>(*this, func_attr);

  if (!func) {
    return emitError("'f' attribute refers to an undefined function: ")

    return false;
  }

  ModuleOp module = op->getParentOfType<ModuleOp>();
  func::FuncOp func = module.lookupSymbol<func::FuncOp>(
      op.getShapeInferenceGraphAttr().getRootReference());

  if (func == nullptr) return false;

  std::vector<_XlaRecvAtHostOp> xla_recv_at_host_ops;
  func.walk([&](_XlaRecvAtHostOp op) { xla_recv_at_host_ops.push_back(op); });