code: def allreduce_grads(all_grads, average):
      
  
      if get_tf_version_tuple() <= (1, 12):
          from tensorflow.contrib import nccl
      else:
          from tensorflow.python.ops import nccl_ops as nccl
      nr_tower = len(all_grads)
      if nr_tower == 1:
          return all_grads
      new_all_grads = []  # N x K
      for grads in zip(*all_grads):
          summed = nccl.all_sum(grads)
  
          grads_for_devices = []  # K
          for g in summed:
              with tf.device(g.device):
                  # tensorflow/benchmarks didn't average gradients
                  if average:
                      g = tf.multiply(g, 1.0 / nr_tower)
              grads_for_devices.append(g)
          new_all_grads.append(grads_for_devices)
  
      # transpose to K x N
      ret = list(zip(*new_all_grads))
      return ret
docs: All-reduce average the gradients among K devices. Results are broadcasted to all devices.
  
      Args:
          all_grads (K x N): List of list of gradients. N is the number of variables.
          average (bool): average gradients or not.
  
      Returns:
          K x N: same as input, but each grad is replaced by the average over K devices.