code: def _bn_agg_mean_var(
      x: jnp.ndarray,
      axis: Union[Tuple[int, ...], int],
      p_agg: bool,
      *,
      axis_name: Optional[str] = None,
      axis_index_groups: Optional[Sequence[Sequence[int]]] = None,
      spatial_shape: Optional[jnp.ndarray] = None,
  ):
      """Aggregate batch statistics accross devices.
  
      Args:
        x: Inputs to compute batch statistics on.
        axis: Reduction axes for the stats.
        p_agg: If True, parallel aggregation is performed using psum.
        axis_name: Name of the axis for psum aggregation.
        axis_index_groups: Groups of axis indices within that named axis
          representing subsets of devices to reduce over (default: None).
        spatial_shape: Per-input spatial shape of *unpadded* input data with
            dimensions (batch, spatial_dims, channels).
  
      Returns:
        Batch stats (mean and variance).
      """
      # When using spatial padding, we cannot accumulate the mean directly and
      # instead must aaccumulate the numerator and sum directly.
      acc_sum = jnp.sum(x, axis=axis, keepdims=False)
      acc_sum2 = jnp.sum(lax.square(x), axis=axis, keepdims=False)
      if spatial_shape is None:
          denom = np.prod([x.shape[i] for i in axis])
      else:
          reduction_axis_shifted = tuple(i - 1 for i in axis if i > 0)
          denom = jnp.prod(
              spatial_shape[:, reduction_axis_shifted], axis=-1
          )
          denom = jnp.sum(denom)
  
      if p_agg:
          concatenated_acc_sum = jnp.concatenate(
              [acc_sum, acc_sum2, denom]
          )
          acc_sum, acc_sum2, denom = jnp.split(
              lax.psum(
                  concatenated_acc_sum,
                  axis_name=axis_name,
                  axis_index_groups=axis_index_groups,
              ),
              3,
          )
  
      denom = jnp.maximum(denom, 1.0)
      mean = acc_sum / denom
      var = acc_sum2 / denom - lax.square(mean)
      return mean, var
  
imports: import functools
  from typing import Optional, Tuple, Union, Callable, Any, Sequence, List, Iterable
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: _bn_agg_mean_var(x, axis, p_agg) and returns: mean,  var