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

    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.
    
tensorpack/graph_builder/utils.py