Aggregate gradients, controlling device for the aggregation.

  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over towers. The inner list is over individual gradients.
    use_mean: if True, mean is taken, else sum of gradients is taken.
    check_inf_nan: If true, check grads for nans and infs.

  Returns:
    The tuple ([(average_gradient, variable),], has_nan_or_inf) where the
      gradient has been averaged across all towers. The variable is chosen from
      the first tower. The has_nan_or_inf indicates the grads has nan or inf.

def aggregate_gradients_using_copy_with_device_selection(
        tower_grads, avail_devices, use_mean=True, check_inf_nan=False):
    """Aggregate gradients, controlling device for the aggregation.

  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over towers. The inner list is over individual gradients.
    use_mean: if True, mean is taken, else sum of gradients is taken.
    check_inf_nan: If true, check grads for nans and infs.

  Returns:
    The tuple ([(average_gradient, variable),], has_nan_or_inf) where the
      gradient has been averaged across all towers. The variable is chosen from
      the first tower. The has_nan_or_inf indicates the grads has nan or inf.
  """
    agg_grads = []
    has_nan_or_inf_list = []
    for i, single_grads in enumerate(zip(*tower_grads)):
        with tf.device(avail_devices[i % len(avail_devices)]):
            grad_and_var, has_nan_or_inf = aggregate_single_gradient(
                single_grads, use_mean, check_inf_nan)
            agg_grads.append(grad_and_var)
            has_nan_or_inf_list.append(has_nan_or_inf)
    return agg_grads