// attribute.
  for (auto &pcall_op : outermost_pcall_ops) {
    auto call = llvm::cast<CallOpInterface>(pcall_op.getOperation());
    CallInterfaceCallable callable = call.getCallableForCallee();
    auto sym = callable.get<mlir::SymbolRefAttr>();
    EncapsulatePartitionedCall(pcall_op, sym.getRootReference());
  }
  // Partitioned calls are executed asynchronous. The calls outside of

    //
    // will be transformed to
    //
    // func.func @f(%x : bar) {
    //   ...
    // }
    // ...
    // func.call @f(%x : bar)

    CallInterfaceCallable callable = call.getCallableForCallee();
    mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>();
    auto symbol =
        mlir::SymbolTable::lookupNearestSymbolFrom(call.getOperation(), sym);
    if (!symbol) return failure();

    args_to_remap.insert(std::make_pair(op, return_to_operand));
  });

  // Find all callers
  module->walk([&](CallOpInterface op) {
    auto callable = op.getCallableForCallee();
    mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>();
    if (!sym) return;
    Operation* func = mlir::SymbolTable::lookupNearestSymbolFrom(op, sym);
    if (!args_to_erase.count(func)) return;

    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }

    // Return the callee of this operation.
    CallInterfaceCallable getCallableForCallee() { return getCalleeAttr(); }
    // Sets the callee from the callable
    void setCalleeFromCallable(CallInterfaceCallable callee);
  }];

  let assemblyFormat = [{

    // returns the function that this operation will launch.
    func::FuncOp getFuncOp() {
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, getFuncAttr());
    }
    CallInterfaceCallable getCallableForCallee() {
      return getFuncAttr();
    }
    void setCalleeFromCallable(::mlir::CallInterfaceCallable callee);
  }];
}

def TfDevice_RemoteRunOp : TfDevice_Op<"remote_run",

  }
  for (auto [ret, operand] :
       llvm::zip(call.getResults(), yield.getOperands())) {
    if (ret != operand) return std::nullopt;
  }
  SymbolRefAttr symbol = call.getCallableForCallee().get<SymbolRefAttr>();
  if (!symbol) return std::nullopt;
  return symbol.getLeafReference();
}

// Extracts the contents of a region with a single block into a new function.

    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }

    // Returns the callee of this operation.
    CallInterfaceCallable getCallableForCallee() { return getFAttr(); }
    // Sets the callee from the callable
    void setCalleeFromCallable(::mlir::CallInterfaceCallable callee);

    // Returns the resolved callee function of this operation.

                           StringAttr old_group) {
  // Adds the TPUReplicateMetatdataOp and TPUCompilationResultOp ops to the
  // function called by the provided 'caller'.
  mlir::CallInterfaceCallable callable = caller.getCallableForCallee();
  mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>();
  auto func = llvm::dyn_cast<mlir::func::FuncOp>(
      mlir::SymbolTable::lookupNearestSymbolFrom(caller, sym));

                           StringAttr old_group) {
  // Adds the TPUReplicateMetatdataOp and TPUCompilationResultOp ops to the
  // function called by the provided 'caller'.
  mlir::CallInterfaceCallable callable = caller.getCallableForCallee();
  mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>();
  auto func = llvm::dyn_cast<mlir::func::FuncOp>(
      mlir::SymbolTable::lookupNearestSymbolFrom(caller, sym));

    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }

    // Returns the callee of this operation.
    CallInterfaceCallable getCallableForCallee() { return getFAttr(); }
    // Sets the callee from the callable.
    void setCalleeFromCallable(CallInterfaceCallable callee);

    // returns the callee of this operation.
    func::FuncOp func() {