code: def _compute_stats(
      x: Array,
      axes: Axes,
      dtype: Optional[Dtype],
      axis_name: Optional[str] = None,
      axis_index_groups: Any = None,
      use_mean: bool = True,
      use_fast_variance: bool = True,
      mask: Optional[Array] = None,
  ):
      """Computes mean and variance statistics.
  
      This implementation takes care of a few important details:
      - Computes in float32 precision for stability in half precision training.
      - If `use_fast_variance` is `True`, mean and variance are computed using
        Var = E[|x|^2] - |E[x]|^2, instead of Var = E[|x - E[x]|^2]), in a single
        XLA fusion.
      - 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: Optional dtype specifying the minimal precision. Statistics are
          always at least float32 for stability (default: dtype of x).
        axis_name: Optional name for the pmapped axis to compute mean over. Note,
          this is only used for pmap and shard map. For SPMD jit, you do not need to
          manually synchronize. Just make sure that the axes are correctly annotated
          and XLA:SPMD will insert the necessary collectives.
        axis_index_groups: 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.
        use_fast_variance: If true, use a faster, but less numerically stable,
          calculation for the variance.
        mask: Binary array of shape broadcastable to `inputs` tensor, indicating
          the positions for which the mean and variance should be computed.
  
      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)
      axes = _canonicalize_axes(x.ndim, axes)
  
      def maybe_distributed_mean(*xs, mask=None):
          mus = tuple(x.mean(axes, where=mask) for x in xs)
          if axis_name is None:
              return mus if len(xs) > 1 else mus[0]
          else:
              # In the distributed case we stack multiple arrays to speed comms.
              if len(xs) > 1:
                  reduced_mus = lax.pmean(
                      jnp.stack(mus, axis=0),
                      axis_name,
                      axis_index_groups=axis_index_groups,
                  )
                  return tuple(reduced_mus[i] for i in range(len(xs)))
              else:
                  return lax.pmean(
                      mus[0],
                      axis_name,
                      axis_index_groups=axis_index_groups,
                  )
  
      if use_mean:
          if use_fast_variance:
              mu, mu2 = maybe_distributed_mean(x, _abs_sq(x), mask=mask)
              # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due
              # to floating point round-off errors.
              var = jnp.maximum(0.0, mu2 - _abs_sq(mu))
          else:
              mu = maybe_distributed_mean(x, mask=mask)
              var = maybe_distributed_mean(
                  _abs_sq(x - jnp.expand_dims(mu, axes)), mask=mask
              )
      else:
          var = maybe_distributed_mean(_abs_sq(x), mask=mask)
          mu = jnp.zeros_like(var)
      return mu, var
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  from jax import lax
  from jax.nn import initializers
  from flax.linen import dtypes, module, transforms
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, use_fast_variance, mask) and returns: mu,  var, unknown, tuple, lax