def _psum(tensor, axis_name=None):
  """Sum all-reduction.

  Args:
    tensor: A tensor.
    axis_name: The axis name to reduce. Must equal to that of the surrounding
      pmap.

  Returns:
    The sum of the `tensor` replicas on each participating devices.
  """
  if axis_name != _pmap_config.axis_name():
    raise ValueError("axis_name (%s) is not equal to that of the surrounding "
                     "pmap (%s)" % (axis_name, _pmap_config.axis_name()))
  devices = _pmap_config.devices()
  if devices is None:
    raise ValueError("Can't retrieve the device list from the surrounding pmap")
  tensor = tf_np.asarray(tensor)
  if tpu_devices(devices):
    # TODO(b/170895907): Remove this workaround when tpu.cross_replica_sum
    #   supports int64/float64.
    is_int64 = False
    is_float64 = False
    if tensor.dtype == np.int64:
      is_int64 = True
      tensor = tensor.astype(np.int32)
    elif tensor.dtype == np.float64:
      is_float64 = True
      tensor = tensor.astype(np.float32)
    # TODO(wangpeng): Supply the `group_assignment` argument to
    #   tpu.cross_replica_sum, calculated from `devices`.
    tensor = tf.compat.v1.tpu.cross_replica_sum(tensor)
    if is_int64:
      tensor = tf.cast(tensor, tf.int64)
    elif is_float64:
      tensor = tf.cast(tensor, tf.float64)
  else:
    tensor = tf.raw_ops.CollectiveReduce(
        input=tensor,
        group_size=len(devices),
        group_key=_GROUP_KEY,
        instance_key=_get_instance_key(),
        merge_op="Add",
        final_op="Id",
        subdiv_offsets=(0,))
  return tf_np.asarray(tensor)
Sum all-reduction.

  Args:
    tensor: A tensor.
    axis_name: The axis name to reduce. Must equal to that of the surrounding
      pmap.

  Returns:
    The sum of the `tensor` replicas on each participating devices.
  
trax/tf_numpy/extensions/extensions.py