class_docstr: Splits the first dimension ("channels") into groups of fixed size. Computes a mean
  and standard deviation over the contents of each group, and uses these to normalise
  the group. Optionally applies a channel-wise affine transformation afterwards.
  
  Given an input array $x$ of shape `(channels, ...)`, this layer splits this up into
  `groups`-many arrays $x_i$ each of shape `(channels/groups, ...)`, and for each one
  computes
  
  $$\frac{x_i - \mathbb{E}[x_i]}{\sqrt{\text{Var}[x_i] + \varepsilon}} * \gamma_i + \beta_i$$
  
  where $\gamma_i$, $\beta_i$ have shape `(channels/groups,)` if
  `channelwise_affine=True`, and $\gamma = 1$, $\beta = 0$ if
  `channelwise_affine=False`.
  
  ??? cite
      [Group Normalization](https://arxiv.org/abs/1803.08494)
  
      ```bibtex
      @article{wu2018group,
          author={Yuxin Wu and Kaiming He},
          title={Group Normalization},
          year={2018},
          journal={arXiv:1803.08494},
      }
      ```
code: def __call__(
      self,
      x: Array,
      state: State,
      *,
      key: Optional[PRNGKeyArray] = None,
  ) -> tuple[Array, State]:
      ...
  
imports: import functools
  import warnings
  from collections.abc import Sequence
  from typing import Optional, overload, Union
  import jax
  import jax.numpy
  from jaxtyping import Array, Float, PRNGKeyArray
  from .._custom_types import sentinel
  from .._misc import left_broadcast_to
  from .._module import field, Module
  from ._stateful import State
label_desc: Use image node_id 3 for calling the GroupNorm obj's underlying member method code with example usage: obj.__call__(x, state) without return types