brainstate 0.1.0.post20241122__py2.py3-none-any.whl → 0.1.0.post20241129__py2.py3-none-any.whl

brainstate/nn/_interaction/_normalizations.py

@@ -17,8 +17,7 @@
 
 from __future__ import annotations
 
-import numbers
-from typing import Callable, Union, Sequence, Optional, Any
+from typing import Callable, Union, Sequence, Optional, Any, Dict
 
 import jax
 import jax.numpy as jnp
@@ -29,10 +28,50 @@ from brainstate.nn._module import Module
 from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
 
 __all__ = [
-    'BatchNorm0d', 'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d',
+    'BatchNorm0d',
+    'BatchNorm1d',
+    'BatchNorm2d',
+    'BatchNorm3d',
+    'LayerNorm',
+    'RMSNorm',
+    'GroupNorm',
 ]
 
 
+def canonicalize_dtype(
+    *args,
+    dtype: jax.typing.DTypeLike | None = None,
+    inexact: bool = True
+) -> jax.typing.DTypeLike:
+    """Canonicalize an optional dtype to the definitive dtype.
+
+    If the ``dtype`` is None this function will infer the dtype. If it is not
+    None it will be returned unmodified or an exceptions is raised if the dtype
+    is invalid.
+    from the input arguments using ``jnp.result_type``.
+
+    Args:
+      *args: JAX array compatible values. None values
+        are ignored.
+      dtype: Optional dtype override. If specified the arguments are cast to
+        the specified dtype instead and dtype inference is disabled.
+      inexact: When True, the output dtype must be a subdtype
+      of `jnp.inexact`. Inexact dtypes are real or complex floating points. This
+      is useful when you want to apply operations that don't work directly on
+      integers like taking a mean for example.
+    Returns:
+      The dtype that *args should be cast to.
+    """
+    if dtype is None:
+        args_filtered = [jnp.asarray(x) for x in args if x is not None]
+        dtype = jnp.result_type(*args_filtered)
+        if inexact and not jnp.issubdtype(dtype, jnp.inexact):
+            dtype = jnp.promote_types(jnp.float32, dtype)
+    if inexact and not jnp.issubdtype(dtype, jnp.inexact):
+        raise ValueError(f'Dtype must be inexact: {dtype}')
+    return dtype
+
+
 def _canonicalize_axes(ndim: int, feature_axes: Sequence[int]):
     axes = []
     for axis in feature_axes:
@@ -59,28 +98,38 @@ def _compute_stats(
     axis_name: Optional[str] = None,
     axis_index_groups: Optional[Sequence[int]] = None,
     use_mean: bool = True,
+    use_fast_variance: bool = True,
+    mask: Optional[jax.Array] = None,
 ):
-    """Computes mean and variance statistics.
+    """
+    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]).
+    - 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: 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.
+        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: 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, val)``.
@@ -89,32 +138,38 @@ def _compute_stats(
         dtype = jax.numpy.result_type(x)
     # promote x to at least float32, this avoids half precision computation
     # but preserves double or complex floating points
-    dtype = jax.numpy.promote_types(dtype, environ.dftype())
+    dtype = jax.numpy.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 = jax.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 jax.lax.pmean(mus[0], axis_name, axis_index_groups=axis_index_groups)
 
-    # Compute mean and mean of squared values.
-    mean2 = jnp.mean(_abs_sq(x), axes)
     if use_mean:
-        mean = jnp.mean(x, axes)
+        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:
-        mean = jnp.zeros(mean2.shape, dtype=dtype)
-
-    # If axis_name is provided, we need to average the mean and mean2 across
-    if axis_name is not None:
-        concatenated_mean = jnp.concatenate([mean, mean2])
-        mean, mean2 = jnp.split(
-            jax.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
+        var = maybe_distributed_mean(_abs_sq(x), mask=mask)
+        mu = jnp.zeros_like(var)
+    return mu, var
 
 
 def _normalize(
@@ -122,9 +177,10 @@ def _normalize(
     mean: Optional[ArrayLike],
     var: Optional[ArrayLike],
     weights: Optional[ParamState],
-    reduction_axes: Sequence[int],
+    reduction_axes: Axes,
+    feature_axes: Axes,
     dtype: DTypeLike,
-    epsilon: Union[numbers.Number, jax.Array],
+    epsilon: jax.typing.ArrayLike,
 ):
     """Normalizes the input of a normalization layer and optionally applies a learned scale and bias.
 
@@ -134,6 +190,7 @@ def _normalize(
       var: Variance to use for normalization.
       weights: The scale and bias parameters.
       reduction_axes: The axes in ``x`` to reduce.
+      feature_axes: The feature axes to apply the scale and bias.
       dtype: The dtype of the result (default: infer from input and params).
       epsilon: Normalization epsilon.
 
@@ -142,16 +199,23 @@ def _normalize(
     """
     if mean is not None:
         assert var is not None, 'mean and val must be both None or not None.'
+        reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)
+        feature_axes = _canonicalize_axes(x.ndim, feature_axes)
         stats_shape = list(x.shape)
         for axis in reduction_axes:
             stats_shape[axis] = 1
         mean = mean.reshape(stats_shape)
         var = var.reshape(stats_shape)
+        feature_shape = [1] * x.ndim
+        for ax in feature_axes:
+            feature_shape[ax] = x.shape[ax]
         y = x - mean
-        mul = jax.lax.rsqrt(var + jnp.asarray(epsilon, dtype))
+        mul = jax.lax.rsqrt(var + epsilon)
         y = y * mul
+        args = []
         if weights is not None:
-            y = _scale_operation(y, weights.value)
+            y, args = _scale_operation(y, weights.value)
+        dtype = canonicalize_dtype(x, *args, dtype=dtype)
     else:
         assert var is None, 'mean and val must be both None or not None.'
         assert weights is None, 'scale and bias are not supported without mean and val'
@@ -159,12 +223,15 @@ def _normalize(
     return jnp.asarray(y, dtype)
 
 
-def _scale_operation(x, param):
+def _scale_operation(x: jax.Array, param: Dict):
+    args = []
     if 'scale' in param:
         x = x * param['scale']
+        args.append(param['scale'])
     if 'bias' in param:
         x = x + param['bias']
-    return x
+        args.append(param['bias'])
+    return x, args
 
 
 class _BatchNorm(Module):
@@ -175,6 +242,7 @@ class _BatchNorm(Module):
         self,
         in_size: Size,
         feature_axis: Axes = -1,
+        *,
         track_running_stats: bool = True,
         epsilon: float = 1e-5,
         momentum: float = 0.99,
@@ -183,14 +251,15 @@ class _BatchNorm(Module):
         scale_initializer: Union[ArrayLike, Callable] = init.Constant(1.),
         axis_name: Optional[Union[str, Sequence[str]]] = None,
         axis_index_groups: Optional[Sequence[Sequence[int]]] = None,
+        use_fast_variance: bool = True,
         name: Optional[str] = None,
         dtype: Any = None,
     ):
         super().__init__(name=name)
 
         # parameters
-        self.in_size = tuple(in_size)
-        self.out_size = tuple(in_size)
+        self.in_size = in_size
+        self.out_size = in_size
         self.affine = affine
         self.bias_initializer = bias_initializer
         self.scale_initializer = scale_initializer
@@ -198,15 +267,17 @@ class _BatchNorm(Module):
         self.track_running_stats = track_running_stats
         self.momentum = jnp.asarray(momentum, dtype=self.dtype)
         self.epsilon = jnp.asarray(epsilon, dtype=self.dtype)
+        self.use_fast_variance = use_fast_variance
 
         # parameters about axis
         feature_axis = (feature_axis,) if isinstance(feature_axis, int) else feature_axis
-        self.feature_axis = _canonicalize_axes(len(in_size), feature_axis)
+        self.feature_axes = _canonicalize_axes(len(self.in_size), feature_axis)
         self.axis_name = axis_name
         self.axis_index_groups = axis_index_groups
 
         # variables
-        feature_shape = tuple([ax if i in self.feature_axis else 1 for i, ax in enumerate(in_size)])
+        feature_shape = tuple([(ax if i in self.feature_axes else 1)
+                               for i, ax in enumerate(self.in_size)])
         if self.track_running_stats:
             self.running_mean = LongTermState(jnp.zeros(feature_shape, dtype=self.dtype))
             self.running_var = LongTermState(jnp.ones(feature_shape, dtype=self.dtype))
@@ -223,7 +294,7 @@ class _BatchNorm(Module):
         else:
             self.weight = None
 
-    def update(self, x):
+    def update(self, x, mask: Optional[jax.Array] = None):
         # input shape and batch mode or not
         if x.ndim == self.num_spatial_dims + 2:
             x_shape = x.shape[1:]
@@ -239,9 +310,9 @@ class _BatchNorm(Module):
 
         # reduce the feature axis
         if batch:
-            reduction_axes = tuple(i for i in range(x.ndim) if (i - 1) not in self.feature_axis)
+            reduction_axes = tuple(i for i in range(x.ndim) if (i - 1) not in self.feature_axes)
         else:
-            reduction_axes = tuple(i for i in range(x.ndim) if i not in self.feature_axis)
+            reduction_axes = tuple(i for i in range(x.ndim) if i not in self.feature_axes)
 
         # fitting phase
         fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
@@ -255,6 +326,8 @@ class _BatchNorm(Module):
                     dtype=self.dtype,
                     axis_name=self.axis_name,
                     axis_index_groups=self.axis_index_groups,
+                    use_fast_variance=self.use_fast_variance,
+                    mask=mask,
                 )
                 self.running_mean.value = self.momentum * self.running_mean.value + (1 - self.momentum) * mean
                 self.running_var.value = self.momentum * self.running_var.value + (1 - self.momentum) * var
@@ -265,14 +338,22 @@ class _BatchNorm(Module):
             mean, var = None, None
 
         # normalize
-        return _normalize(x, mean, var, self.weight, reduction_axes, self.dtype, self.epsilon)
+        return _normalize(
+            x,
+            mean=mean,
+            var=var,
+            weights=self.weight,
+            reduction_axes=reduction_axes,
+            feature_axes=self.feature_axes,
+            dtype=self.dtype,
+            epsilon=self.epsilon
+        )
 
 
 class BatchNorm0d(_BatchNorm):
-    r"""1-D batch normalization [1]_.
+    r"""0-D batch normalization [1]_.
 
-    The data should be of `(b, l, c)`, where `b` is the batch dimension,
-    `l` is the layer dimension, and `c` is the channel dimension.
+    The data should be of `(b, c)`, where `b` is the batch dimension, and `c` is the channel dimension.
 
     %s
     """
@@ -375,7 +456,10 @@ _bn_doc = r'''
     example, `[[0, 1], [2, 3]]` would independently batch-normalize over
     the examples on the first two and last two devices. See `jax.lax.psum`
     for more details.
-
+  use_fast_variance: If true, use a faster, but less numerically stable,
+    calculation for the variance.
+    
+    
   References
   ----------
   .. [1] Ioffe, Sergey and Christian Szegedy. “Batch Normalization: Accelerating Deep Network Training
@@ -386,3 +470,441 @@ _bn_doc = r'''
 BatchNorm1d.__doc__ = BatchNorm1d.__doc__ % _bn_doc
 BatchNorm2d.__doc__ = BatchNorm2d.__doc__ % _bn_doc
 BatchNorm3d.__doc__ = BatchNorm3d.__doc__ % _bn_doc
+
+
+class LayerNorm(Module):
+    """
+    Layer normalization (https://arxiv.org/abs/1607.06450).
+
+    LayerNorm normalizes the activations of the layer for each given example in a
+    batch independently, rather than across a batch like Batch Normalization.
+    i.e. applies a transformation that maintains the mean activation within
+    each example close to 0 and the activation standard deviation close to 1.
+
+    Example usage::
+
+      >>> import brainstate as bst
+      >>> x = bst.random.normal(size=(3, 4, 5, 6))
+      >>> layer = bst.nn.LayerNorm(x.shape)
+      >>> layer.states()
+      >>> y = layer(x)
+
+    Attributes:
+      in_size: The input shape, without batch size.
+      epsilon: A small float added to variance to avoid dividing by zero.
+      dtype: the dtype of the result (default: infer from input and params).
+      use_bias:  If True, bias (beta) is added.
+      use_scale: If True, multiply by scale (gamma). When the next layer is linear
+          (also e.g. nnx.relu), this can be disabled since the scaling will be done
+          by the next layer.
+      bias_init: Initializer for bias, by default, zero.
+      scale_init: Initializer for scale, by default, one.
+      reduction_axes: Axes for computing normalization statistics. It is recommended
+            to use the negative integer, since when the batch dimension is used,
+            the reduction_axes may be wrong when using the positive integer.
+      feature_axes: Feature axes for learned bias and scaling.
+      axis_name: the axis name used to combine batch statistics from multiple
+          devices. See ``jax.pmap`` for a description of axis names (default: None).
+          This is only needed if the model is subdivided across devices, i.e. the
+          array being normalized is sharded across devices within a pmap.
+      axis_index_groups: groups of axis indices within that named axis
+          representing subsets of devices to reduce over (default: None). For
+          example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over
+          the examples on the first two and last two devices. See ``jax.lax.psum``
+          for more details.
+      use_fast_variance: If true, use a faster, but less numerically stable,
+          calculation for the variance.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        reduction_axes: Axes = -1,
+        feature_axes: Axes = -1,
+        *,
+        epsilon: float = 1e-6,
+        use_bias: bool = True,
+        use_scale: bool = True,
+        bias_init: Callable = init.ZeroInit(),
+        scale_init: Callable = init.Constant(1.0),
+        axis_name: Optional[str] = None,
+        axis_index_groups: Any = None,
+        use_fast_variance: bool = True,
+        dtype: Optional[jax.typing.DTypeLike] = None,
+    ):
+        super().__init__()
+
+        self.in_size = in_size
+        self.out_size = in_size
+
+        # parameters about axis
+        feature_axes = (feature_axes,) if isinstance(feature_axes, int) else feature_axes
+        self.feature_axes = _canonicalize_axes(len(self.in_size), feature_axes)
+        self.reduction_axes = (reduction_axes,) if isinstance(reduction_axes, int) else reduction_axes
+        self.axis_name = axis_name
+        self.axis_index_groups = axis_index_groups
+
+        # variables
+        feature_shape = tuple([(ax if i in self.feature_axes else 1)
+                               for i, ax in enumerate(self.in_size)])
+
+        weights = dict()
+        if use_scale:
+            weights['scale'] = init.param(scale_init, feature_shape)
+        if use_bias:
+            weights['bias'] = init.param(bias_init, feature_shape)
+        if len(weights):
+            self.weight = ParamState(weights)
+        else:
+            self.weight = None
+
+        # parameters
+        self.epsilon = epsilon
+        self.dtype = dtype or environ.dftype()
+        self.use_bias = use_bias
+        self.use_scale = use_scale
+        self.bias_init = bias_init
+        self.scale_init = scale_init
+        self.use_fast_variance = use_fast_variance
+
+    def update(self, x, *, mask: Optional[jax.Array] = None):
+        """Applies layer normalization on the input.
+
+        Args:
+          x: the inputs
+
+        Returns:
+          Normalized inputs (the same shape as inputs).
+        """
+        mean, var = _compute_stats(
+            x,
+            self.reduction_axes,
+            dtype=self.dtype,
+            axis_name=self.axis_name,
+            axis_index_groups=self.axis_index_groups,
+            use_fast_variance=self.use_fast_variance,
+            mask=mask,
+        )
+
+        return _normalize(
+            x,
+            mean=mean,
+            var=var,
+            weights=self.weight,
+            reduction_axes=self.reduction_axes,
+            feature_axes=self.feature_axes,
+            dtype=self.dtype,
+            epsilon=self.epsilon,
+        )
+
+
+class RMSNorm(Module):
+    """
+    RMS Layer normalization (https://arxiv.org/abs/1910.07467).
+
+    RMSNorm normalizes the activations of the layer for each given example in a
+    batch independently, rather than across a batch like Batch Normalization.
+    Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the
+    standard deviation of the activations, RMSNorm does not re-center at all
+    and instead normalizes by the root mean square of the activations.
+
+    Example usage::
+
+      >>> import brainstate as bst
+      >>> x = bst.random.normal(size=(5, 6))
+      >>> layer = bst.nn.RMSNorm(num_features=6)
+      >>> layer.states()
+      >>> y = layer(x)
+
+    Attributes:
+        in_size: The input shape, without batch size.
+        epsilon: A small float added to variance to avoid dividing by zero.
+        dtype: the dtype of the result (default: infer from input and params).
+        use_scale: If True, multiply by scale (gamma). When the next layer is linear
+            (also e.g. nn.relu), this can be disabled since the scaling will be done
+            by the next layer.
+        scale_init: Initializer for scale, by default, one.
+        reduction_axes: Axes for computing normalization statistics. It is recommended
+            to use the negative integer, since when the batch dimension is used,
+            the reduction_axes may be wrong when using the positive integer.
+        feature_axes: Feature axes for learned bias and scaling.
+        axis_name: the axis name used to combine batch statistics from multiple
+            devices. See ``jax.pmap`` for a description of axis names (default: None).
+            This is only needed if the model is subdivided across devices, i.e. the
+            array being normalized is sharded across devices within a pmap.
+        axis_index_groups: groups of axis indices within that named axis
+            representing subsets of devices to reduce over (default: None). For
+            example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over
+            the examples on the first two and last two devices. See ``jax.lax.psum``
+            for more details.
+        use_fast_variance: If true, use a faster, but less numerically stable,
+            calculation for the variance.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        *,
+        epsilon: float = 1e-6,
+        dtype: Optional[jax.typing.DTypeLike] = None,
+        use_scale: bool = True,
+        scale_init: Callable = init.Constant(1.0),
+        reduction_axes: Axes = -1,
+        feature_axes: Axes = -1,
+        axis_name: Optional[str] = None,
+        axis_index_groups: Any = None,
+        use_fast_variance: bool = True,
+    ):
+        super().__init__()
+
+        self.in_size = in_size
+        self.out_size = in_size
+
+        # parameters about axis
+        feature_axes = (feature_axes,) if isinstance(feature_axes, int) else feature_axes
+        self.feature_axes = _canonicalize_axes(len(self.in_size), feature_axes)
+        self.reduction_axes = (reduction_axes, ) if isinstance(reduction_axes, int) else reduction_axes
+        self.axis_name = axis_name
+        self.axis_index_groups = axis_index_groups
+
+        # variables
+        feature_shape = tuple([(ax if i in self.feature_axes else 1)
+                               for i, ax in enumerate(self.in_size)])
+        if use_scale:
+            self.scale = ParamState({'scale': init.param(scale_init, feature_shape)})
+        else:
+            self.scale = None
+
+        # parameters
+        self.epsilon = epsilon
+        self.dtype = dtype or environ.dftype()
+        self.use_scale = use_scale
+        self.scale_init = scale_init
+        self.use_fast_variance = use_fast_variance
+
+    def update(self, x, *, mask: Optional[jax.Array] = None):
+        """Applies layer normalization on the input.
+
+        Args:
+          x: the inputs
+          mask: the mask
+
+        Returns:
+          Normalized inputs (the same shape as inputs).
+        """
+        mean, var = _compute_stats(
+            x,
+            self.reduction_axes,
+            dtype=self.dtype,
+            axis_name=self.axis_name,
+            axis_index_groups=self.axis_index_groups,
+            use_mean=False,
+            use_fast_variance=self.use_fast_variance,
+            mask=mask,
+        )
+
+        return _normalize(
+            x,
+            mean=mean,
+            var=var,
+            weights=self.scale,
+            reduction_axes=self.reduction_axes,
+            feature_axes=self.feature_axes,
+            dtype=self.dtype,
+            epsilon=self.epsilon,
+        )
+
+
+class GroupNorm(Module):
+    """
+    Group normalization (arxiv.org/abs/1803.08494).
+
+    This op is similar to batch normalization, but statistics are shared across
+    equally-sized groups of channels and not shared across batch dimension.
+    Thus, group normalization does not depend on the batch composition and does
+    not require maintaining internal state for storing statistics.
+    The user should either specify the total number of channel groups or the
+    number of channels per group.
+
+    .. note::
+      LayerNorm is a special case of GroupNorm where ``num_groups=1``.
+
+    Example usage::
+
+      >>> import numpy as np
+      >>> import brainstate as bst
+      ...
+      >>> x = bst.random.normal(size=(3, 4, 5, 6))
+      >>> layer = bst.nn.GroupNorm(x.shape, num_groups=3)
+      >>> layer.states()
+      >>> y = layer(x)
+      >>> y = bst.nn.GroupNorm(x.shape, num_groups=1)(x)
+      >>> y2 = bst.nn.LayerNorm(x.shape, reduction_axes=(1, 2, 3))(x)
+      >>> np.testing.assert_allclose(y, y2)
+
+    Attributes:
+      in_size: The input shape, without batch size.
+      num_groups: the total number of channel groups. The default value of 32 is
+        proposed by the original group normalization paper.
+      group_size: the number of channels in a group.
+      epsilon: A small float added to variance to avoid dividing by zero.
+      dtype: the dtype of the result (default: infer from input and params).
+      use_bias:  If True, bias (beta) is added.
+      use_scale: If True, multiply by scale (gamma). When the next layer is linear
+        (also e.g. nn.relu), this can be disabled since the scaling will be done
+        by the next layer.
+      bias_init: Initializer for bias, by default, zero.
+      scale_init: Initializer for scale, by default, one.
+      reduction_axes: List of axes used for computing normalization statistics.
+        This list must include the final dimension, which is assumed to be the
+        feature axis. Furthermore, if the input used at call time has additional
+        leading axes compared to the data used for initialisation, for example due
+        to batching, then the reduction axes need to be defined explicitly.
+        It is recommended to use the negative integer, since when the batch dimension is used,
+        the reduction_axes may be wrong when using the positive integer.
+      axis_name: the axis name used to combine batch statistics from multiple
+        devices. See ``jax.pmap`` for a description of axis names (default: None).
+        This is only needed if the model is subdivided across devices, i.e. the
+        array being normalized is sharded across devices within a pmap or 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: groups of axis indices within that named axis
+        representing subsets of devices to reduce over (default: None). For
+        example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the
+        examples on the first two and last two devices. See ``jax.lax.psum`` for
+        more details.
+      use_fast_variance: If true, use a faster, but less numerically stable,
+        calculation for the variance.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        feature_axis: Axes = -1,
+        num_groups: Optional[int] = 32,
+        group_size: Optional[int] = None,
+        *,
+        epsilon: float = 1e-6,
+        dtype: Optional[jax.typing.DTypeLike] = None,
+        use_bias: bool = True,
+        use_scale: bool = True,
+        bias_init: Callable = init.ZeroInit(),
+        scale_init: Callable = init.Constant(1.),
+        reduction_axes: Optional[Axes] = None,
+        axis_name: Optional[str] = None,
+        axis_index_groups: Any = None,
+        use_fast_variance: bool = True,
+    ):
+        super().__init__()
+
+        self.in_size = in_size
+        self.out_size = in_size
+
+        # parameters about axis
+        feature_axis = (feature_axis,) if isinstance(feature_axis, int) else feature_axis
+        self.feature_axes = _canonicalize_axes(len(self.in_size), feature_axis)
+        self.reduction_axes = (reduction_axes,) if isinstance(reduction_axes, int) else reduction_axes
+        self.axis_name = axis_name
+        self.axis_index_groups = axis_index_groups
+
+        if (num_groups is None and group_size is None) or (
+            num_groups is not None and group_size is not None
+        ):
+            raise ValueError(
+                'Either `num_groups` or `group_size` should be '
+                'specified. If `group_size` is to be specified, '
+                'pass `num_groups=None` as argument to override '
+                'the default `num_groups` value of 32.'
+            )
+
+        feature_shape = tuple([(ax if i in self.feature_axes else 1)
+                               for i, ax in enumerate(self.in_size)])
+        assert len(feature_shape) == 1, 'GroupNorm only supports 1D feature axis.'
+        num_features = feature_shape[0]
+        if group_size is not None:
+            if num_features % group_size != 0:
+                raise ValueError(
+                    'Number of features ({}) is not multiple of the '
+                    'group size ({}).'.format(num_features, group_size)
+                )
+            self.num_groups = num_features // group_size
+            self.group_size = group_size
+        else:
+            if not isinstance(num_groups, int) or num_groups <= 0 or (
+                num_features % num_groups != 0
+            ):
+                raise ValueError(
+                    'Number of groups ({}) does not divide the number'
+                    ' of channels ({}).'.format(num_groups, num_features)
+                )
+            self.num_groups = num_groups
+            self.group_size = num_features // num_groups
+
+        # variables
+        weights = dict()
+        if use_scale:
+            weights['scale'] = init.param(scale_init, feature_shape)
+        if use_bias:
+            weights['bias'] = init.param(bias_init, feature_shape)
+        if len(weights):
+            self.weight = ParamState(weights)
+        else:
+            self.weight = None
+
+        # parameters
+        self.epsilon = epsilon
+        self.dtype = dtype
+        self.use_bias = use_bias
+        self.use_scale = use_scale
+        self.bias_init = bias_init
+        self.scale_init = scale_init
+        self.use_fast_variance = use_fast_variance
+
+    def update(self, x, *, mask: Optional[jax.Array] = None):
+        """Applies group normalization to the input (arxiv.org/abs/1803.08494).
+
+        Args:
+          x: the input of shape ``...self.num_features`` where ``self.num_features``
+            is a channels dimension and ``...`` represents an arbitrary number of
+            extra dimensions that can be used to accumulate statistics over. If no
+            reduction axes have been specified then all additional dimensions ``...``
+            will be used to accumulate statistics apart from the leading dimension
+            which is assumed to represent the batch.
+          mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating
+            the positions for which the mean and variance should be computed.
+
+        Returns:
+          Normalized inputs (the same shape as inputs).
+        """
+        if self.reduction_axes is not None:
+            reduction_axes = self.reduction_axes
+        else:
+            reduction_axes = list(range(1, x.ndim - 1)) + [-1]
+        reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)
+
+        group_shape = x.shape[:-1] + (self.num_groups, self.group_size)
+        if mask is not None:
+            mask = mask.reshape(mask.shape[:-1] + (self.num_groups, self.group_size))
+
+        mean, var = _compute_stats(
+            x.reshape(group_shape),
+            list(reduction_axes[:-1]) + [-1],
+            dtype=self.dtype,
+            axis_name=self.axis_name,
+            axis_index_groups=self.axis_index_groups,
+            use_fast_variance=self.use_fast_variance,
+            mask=mask,
+        )
+        mean = jnp.repeat(mean, self.group_size, axis=1)
+        var = jnp.repeat(var, self.group_size, axis=1)
+        return _normalize(
+            x,
+            mean=mean,
+            var=var,
+            weights=self.weight,
+            reduction_axes=reduction_axes[:-1],
+            feature_axes=self.feature_axes,
+            dtype=self.dtype,
+            epsilon=self.epsilon,
+        )