code: def _compute_stats(
      x: Array,
      axes: tp.Optional[Axes],
      dtype: tp.Optional[Dtype],
      axis_name: tp.Optional[str] = None,
      axis_index_groups: tp.Any = None,
      use_mean: bool = True,
  ):
      """Computes mean and variance statistics.
  
      This implementation takes care of a few important details:
      - Computes in float32 precision for stability in half precision training.
      - mean and variance are computable in a single XLA fusion,
        by using Var = E[|x|^2] - |E[x]|^2 instead of Var = E[|x - E[x]|^2]).
      - Clips negative variances to zero which can happen due to
        roundoff errors. This avoids downstream NaNs.
      - Supports averaging across a parallel axis and subgroups of a parallel axis
        with a single `lax.pmean` call to avoid latency.
  
      Arguments:
        x: Input array.
        axes: The axes in ``x`` to compute mean and variance statistics for.
        dtype: tp.Optional dtype specifying the minimal precision. Statistics
          are always at least float32 for stability (default: dtype of x).
        axis_name: tp.Optional name for the pmapped axis to compute mean over.
        axis_index_groups: tp.Optional axis indices.
        use_mean: If true, calculate the mean from the input and use it when
          computing the variance. If false, set the mean to zero and compute
          the variance without subtracting the mean.
  
      Returns:
        A pair ``(mean, var)``.
      """
      if dtype is None:
          dtype = jnp.result_type(x)
      # promote x to at least float32, this avoids half precision computation
      # but preserves double or complex floating points
      dtype = jnp.promote_types(dtype, jnp.float32)
      x = jnp.asarray(x, dtype)
  
      mean2 = jnp.mean(_abs_sq(x), axes)
      if use_mean:
          mean = jnp.mean(x, axes)
      else:
          mean = jnp.zeros(mean2.shape, dtype=dtype)
  
      if axis_name is not None:
          concatenated_mean = jnp.concatenate([mean, mean2])
          mean, mean2 = jnp.split(
              lax.pmean(
                  concatenated_mean,
                  axis_name=axis_name,
                  axis_index_groups=axis_index_groups,
              ),
              2,
          )
      # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due
      # to floating point round-off errors.
      var = jnp.maximum(0.0, mean2 - _abs_sq(mean))
      return mean, var
  
file_path: flax/flax/experimental/nnx/nnx/nn/normalization.py
imports: import typing
  import jax
  import jax.numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib, rnglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 3 for calling a global function with example usage: _compute_stats(x, axes, dtype, axis_name, axis_index_groups, use_mean) and returns: mean,  var