<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def aggregate_grads(all_grads, colocation=False, devices=None, average=True):
 """ 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. """

    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