// Remove the original parallel_execute.
  parallel_execute_op->dropAllUses();
  parallel_execute.erase();
}

// Replaces TPUExecute with TPUExecuteAndUpdateVariables.
void ReplaceExecute(tf_device::LaunchOp execute_launch,
                    tf_device::LaunchOp merged_execute_launch,
                    const VariableAccessesForTPUExecute& var_access_info) {
  // Replace the uses.