code: def allreduce_grads(all_grads, average, mode="nccl"):
      """
      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.
          mode (str): "nccl", "collective"
  
      Returns:
          K x N: same as input, but each grad is replaced by the average over K devices.
      """
      assert mode in ["nccl", "collective"], mode
  
      nr_tower = len(all_grads)
      if nr_tower == 1:
          return all_grads
      new_all_grads = []  # N x K
      for grads in zip(*all_grads):
          # k grads
          if mode == "nccl":
              if get_tf_version_tuple() <= (1, 12):
                  from tensorflow.contrib import nccl  # deprecated
              else:
                  from tensorflow.python.ops import nccl_ops as nccl
              summed = nccl.all_sum(grads)
          else:
              from tensorflow.python.ops import collective_ops
  
              summed = []
              shared_cnt = _get_shared_cnt()
              for t in grads:
                  with tf.device(t.device):
                      t = collective_ops.all_reduce(
                          t,
                          len(grads),
                          42,  # group key is any fixed integer for a fixed group of devices
                          shared_cnt + 100,
                          "Add",
                          "Id",
                          communication_hint="nccl",
                      )
                      summed.append(t)
  
          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
  
imports: import operator
  from contextlib import contextmanager
  import tensorflow
  import threading
  from ..compat import tfv1
  from ..tfutils.common import get_tf_version_tuple
  from ..tfutils.scope_utils import cached_name_scope, under_name_scope
  from ..tfutils.varreplace import custom_getter_scope
  from ..utils import logger
  from ..utils.argtools import call_only_once
label_desc: Use image node_id 6 for calling a global function with example usage: allreduce_grads(all_grads, average, mode) and returns: ret, all_grads