code: def aggregate_grads_colocate(all_grads, average=True):
      """
      Aggregate the gradients. The aggregation is colocated with the variable.
  
      Args:
          all_grads (K x N x 2): A list of K lists. Each of the list is a list of N (grad, var) tuples.
              The variables have to be shared across the K lists.
          average (bool): do average or sum
      Returns:
          (N x 2): A list of N (grad, var) tuples, where grad is averaged or summed over K.
      """
      nr_tower = len(all_grads)
      if nr_tower == 1:
          return all_grads[0]
  
      def aggregate(grads):
          if average:
              return tf.multiply(tf.add_n(grads), 1.0 / nr_tower)
          else:
              return tf.add_n(grads)
  
      ret = []
      for idx, grad_and_vars in enumerate(zip(*all_grads)):
          # Ngpu * 2
          v = grad_and_vars[0][1]
          grads = [g for (g, _) in grad_and_vars]
          with tf.device(v.device):  # colocate summed grad with var
              grad = aggregate(grads)
          ret.append((grad, v))
      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 8 for calling a global function with example usage: aggregate_grads_colocate(all_grads, average) and returns: ret, all_grads, tf, tf