(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.
Computes the mean of a distributed value ``x``.
Computes the mean of a distributed value x .
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)
def mean_or_pmean ( n_devices , x , axis = None ) : 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 )