code: def aggregate_grads(all_grads,
                      colocation=False,
                      devices=None,
                      average=True):
      
      assert not (devices is not None and colocation)
      if devices is not None:
          assert isinstance(devices, list), devices
  
      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]
  
          if colocation:
              with tf.device(v.device):       # colocate summed grad with var
                  grad = aggregate(grads)
          elif devices is None:
              grad = aggregate(grads)
          else:
              dev = devices[idx % len(devices)]
              with tf.device(dev):
                  grad = aggregate(grads)
          ret.append((grad, v))
      return ret
docs: Average the gradients.
  
      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 the same across the K lists.
          colocation (bool): colocate gradient averaging on the device of the variable.
          devices (list[str]): assign the averaging to these device in
              round-robin. Cannot be used together with ``colocation``.
          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.