while_node->ClearAttr("cond");
    while_node->AddAttr("cond", *loop_cond_func);
  }

  return absl::OkStatus();
}

// Rewrites while loop cond function for host.
Status RewriteHostWhileLoopCond(
    const string& cond_host_func_name, const string& while_node_name,
    const string& host_transfer_key, const string& xla_cluster_attr_name,