code: def reduce_by_device(parallelism, data, reduce_fn):
      """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.
      """
      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
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 20 for calling a global function with example usage: reduce_by_device(parallelism, data, reduce_fn) and returns: device_parallelism,  device_parallelism