code: def expand_by_device(original_parallelism, device_parallelism, data):
    
    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]
docs: 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