code: def _n_weights_per_core(weights):  # pylint: disable=invalid-name
      """Calculates the number of weights per core.
  
      In multi-device settings, gradients and losses are averaged over all devices.
      When loss is weighted and the number of weights can differ by device, e.g.,
      when the weights represent the number of tokens in a batch of sentences (which
      can differ from device to device), we want to make sure each token on each
      device is weighted in the same way. This function ensures that by reporting
      the number of weights per core in multi-core settings (and simply
      np.sum(weights) in a single-core setting).
  
      Args:
        weights: tensor with arbitrary shape
  
      Returns:
        a scalar equal to np.sum(weights) in 1-machine settings and to the sum
        of weights over all cores divided by the number of cores otherwise
      """
      weights_sum = jnp.sum(weights)
      if fastmath.global_device_count() < 2:
          return weights_sum
      else:
          try:
              n_devices_total = fastmath.psum(1, "batch")
              return (
                  fastmath.psum(weights_sum, "batch") / n_devices_total
              )
          except (
              NameError,
              ValueError,
          ):  # running outside of pmap, e.g., on init
              return weights_sum
  
imports: from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
label_desc: Use image node_id 2 for calling a global function with example usage: _n_weights_per_core(weights) and returns: weights_sum, unknown, weights_sum