code: def reduce_by_device(parallelism, data, reduce_fn):
    
    unique_devices = []
    device_to_data = {}
    for dev, datum in zip(parallelism.devices, data):
      if dev not in device_to_data:
        unique_devices.append(dev)
        device_to_data[dev] = [datum]
      else:
        device_to_data[dev].append(datum)
    device_parallelism = Parallelism(unique_devices)
    grouped_data = [device_to_data[dev] for dev in unique_devices]
    return device_parallelism, device_parallelism(reduce_fn, grouped_data)
docs: Reduces data per device.
  
    This can be useful, for example, if we want to all-reduce n tensors on k<n
    devices (like during eval when we have only one device).  We call
    reduce_by_device() to first sum the tensors per device, then call our usual
    all-reduce operation to create one sum per device, followed by
    expand_by_device, to create the appropriate number of pointers to these
    results.  See all_reduce_ring() below for an example of how this is used.
  
    Args:
      parallelism: a expert_utils.Parallelism object
      data: a list of Tensors with length parallelism.n
      reduce_fn: a function taking a list of Tensors.  e.g. tf.add_n
  
    Returns:
      device_parallelism: a Parallelism object with each device listed only once.
      reduced_data: A list of Tensors, one per device.