parameters: (replica_grads)
return_statement: return agg_all_g_and_v
docstring: Aggregate gradients using nccl allreduce.
docstring_summary: Aggregate gradients using nccl allreduce.
docstring_tokens: Aggregate gradients using nccl allreduce .
function: def aggregate_gradients_using_nccl(replica_grads):
    """Aggregate gradients using nccl allreduce."""
    agg_all_g_and_v = []
    for single_g_and_v in zip(*replica_grads):
      single_grads = [g for g, _ in single_g_and_v]
      agg_grads = nccl_ops.all_sum(single_grads)
      agg_all_g_and_v.append(
          [(g, v) for g, (_, v) in zip(agg_grads, single_g_and_v)])
  
    agg_all_g_and_v = list(zip(*agg_all_g_and_v))
  
    return agg_all_g_and_v
function_tokens: def aggregate_gradients_using_nccl ( replica_grads ) : agg_all_g_and_v = [ ] for single_g_and_v in zip ( * replica_grads ) : single_grads = [ g for g , _ in single_g_and_v ] agg_grads = nccl_ops . all_sum ( single_grads ) agg_all_g_and_v . append ( [ ( g , v ) for g , ( _ , v ) in zip ( agg_grads , single_g_and_v ) ] ) agg_all_g_and_v = list ( zip ( * agg_all_g_and_v ) ) return agg_all_g_and_v