def mean_or_pmean(n_devices, x, axis=None):
  """Computes the mean of a distributed value ``x``.

  Args:
    n_devices: Number of devices.
    x: Distributed array.
    axis: Axis along which to compute means; can only be ``0`` or ``None``.

  Returns:
    A local array.
  """
  if fastmath.backend_name() == 'tensorflow-numpy' and n_devices > 1:
    if axis not in (None, 0):
      raise ValueError('axis can only be None or 0')
    x = fastmath.pmap(fastmath.psum)(x)[0] / n_devices
    if axis is None:
      x = jnp.mean(x)
    return x
  else:
    return jnp.mean(x, axis=axis)
Computes the mean of a distributed value ``x``.

  Args:
    n_devices: Number of devices.
    x: Distributed array.
    axis: Axis along which to compute means; can only be ``0`` or ``None``.

  Returns:
    A local array.
  
trax/layers/acceleration.py