<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def allreduce_grads(all_grads, average):
 """ 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. Returns: K x N: same as input, but each grad is replaced by the average over K devices. """


    if get_tf_version_tuple() <= (1, 12):
        from tensorflow.contrib import nccl
    else:
        from tensorflow.python.ops import nccl_ops as nccl
    nr_tower = len(all_grads)
    if nr_tower == 1:
        return all_grads
    new_all_grads = []  # N x K
    for grads in zip(*all_grads):
        summed = nccl.all_sum(grads)

        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