Compute the reduction of all Tensors and put the result everywhere.

  Performance-optimized for a ring of devices.

  Args:
    xs: a list of n tf.Tensors
    devices: a list of strings
    reduction_fn_string: "SUM" or "MAX"

  Returns:
    a list of n Tensors
  Raises:
    ValueError: if devices is not a list of n strings

def allreduce_ring(xs, devices, reduction_fn_string="SUM"):
  """Compute the reduction of all Tensors and put the result everywhere.

  Performance-optimized for a ring of devices.

  Args:
    xs: a list of n tf.Tensors
    devices: a list of strings
    reduction_fn_string: "SUM" or "MAX"

  Returns:
    a list of n Tensors
  Raises:
    ValueError: if devices is not a list of n strings
  """
  n = len(xs)
  if len(devices) != n:
    raise ValueError("devices must be a list of length len(xs)")
  if n == 1:
    return xs
  shape = xs[0].shape.as_list()
  # tf.logging.info("allreduce_ring shape = %s" % shape)
  size = None if None in shape else mtf.list_product(shape)
  if size is None or size < 1024 or size % n != 0:
    return allreduce_ring_single_shard(xs, devices, reduction_fn_string)

  def _circular_shift(l, n):
    n %= len(l)
    return l[-n:] + l[:-n]
  def _flatten_and_split(x):
    # tf.reshape treats [-1] as a special value denoting 1D flattening.
    return tf.split(tf.reshape(x, [-1]), n)
  def _concat_and_reshape(xs):
    return tf.reshape(tf.concat(xs, 0), shape)

  # [device, shard]
  x_split = mtf.parallel(devices, _flatten_and_split, xs)
  x_split_t = mtf.transpose_list_of_lists(x_split)

  y_split_t = []
  for shard in xrange(n):
    shard_xs = _circular_shift(x_split_t[shard], shard)
    shard_devices = _circular_shift(devices, shard)
    shard_ys = allreduce_ring_single_shard(
        shard_xs, shard_devices, reduction_fn_string)
    y_split_t.append(_circular_shift(shard_ys, -shard))
  y_split = mtf.transpose_list_of_lists(y_split_t)
  ys = mtf.parallel(devices, _concat_and_reshape, y_split)
  return ys