Opposite of reduce_by_device().

  Args:
    original_parallelism: a expert_utils.Parallelism object.
    device_parallelism: a expert_utils.Parallelism object.
    data: a list of tensors with length device_parallelism.n

  Returns:
    a list of Tensors with length original_parallelism.n

def expand_by_device(original_parallelism, device_parallelism, data):
  """Opposite of reduce_by_device().

  Args:
    original_parallelism: a expert_utils.Parallelism object.
    device_parallelism: a expert_utils.Parallelism object.
    data: a list of tensors with length device_parallelism.n

  Returns:
    a list of Tensors with length original_parallelism.n
  """
  device_to_datum = {
      device_parallelism.devices[i]: data[i]
      for i in range(device_parallelism.n)}
  return [device_to_datum[d] for d in original_parallelism.devices]