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_single_shard(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)
  binary_reduction = mtf.binary_reduction_fn(reduction_fn_string)
  assert len(devices) == n, "devices must be a list of length len(xs)"
  if n == 1:
    return xs
  result = [None] * n
  if n % 2 == 0:
    left_center = n // 2 - 1
    right_center = left_center + 1
  else:
    left_center = n // 2
    right_center = left_center
  left_sum = xs[0]
  for i in xrange(1, left_center + 1):
    with tf.device(devices[i]):
      left_sum = binary_reduction(left_sum, xs[i])
  right_sum = xs[n-1]
  for i in reversed(xrange(left_center + 1, n - 1)):
    with tf.device(devices[i]):
      right_sum = binary_reduction(xs[i], right_sum)
  with tf.device(devices[left_center]):
    result[left_center] = binary_reduction(left_sum, right_sum)
  if n % 2 == 0:
    with tf.device(devices[right_center]):
      result[right_center] = binary_reduction(left_sum, right_sum)
  for i in reversed(xrange(left_center)):
    with tf.device(devices[i]):
      result[i] = tf.identity(result[i + 1])
  for i in xrange(right_center + 1, n):
    with tf.device(devices[i]):
      result[i] = tf.identity(result[i - 1])
  return result