// TODO(bmzhao): After discussing with Allen, we should execute this via a
  // PartitionedCallOp for compatibility with "tooling that assumes functions in
  // graphs are PartitionedCallOps".
  TF_RETURN_IF_ERROR((*out)->Reset(name_.c_str(), nullptr));

  // Adding the user-provided inputs to the function.
  TF_RETURN_IF_ERROR((*out)->AddInputList(inputs));

  // Replace uses.
  op.replaceAllUsesWith(new_op.getResults().take_front(op.getNumResults()));
  op.erase();
  return success();
}

// A resource-lifted function for (potentially multiple) PartitionedCallOps and
// information about the lifting changes.
struct PartitionedCallLiftingInfo {
  // Function with resources lifted. Can be nullptr if nothing needs to change.
  func::FuncOp lifted_callee;