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},
      }
      ```
docstr: **Arguments:**
  
  - `groups`: The number of groups to split the input into.
  - `channels`: The number of input channels. May be left unspecified (e.g. just
      `None`) if `channelwise_affine=False`.
  - `eps`: Value added to denominator for numerical stability.
  - `channelwise_affine`: Whether the module has learnable affine parameters.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","groups","`",":","The","number","of","groups","to","split","the","input","into",".","-","`","channels","`",":","The","number","of","input","channels",".","May","be","left","unspecified","(","e.g",".","just","`","None","`",")","if","`","channelwise_affine=False","`",".","-","`","eps","`",":","Value","added","to","denominator","for","numerical","stability",".","-","`","channelwise_affine","`",":","Whether","the","module","has","learnable","affine","parameters","."]
code: def __init__(
      self,
      groups: int,
      channels: Optional[int] = None,
      eps: float = 1e-5,
      channelwise_affine: bool = True,
  ):
      """**Arguments:**
  
      - `groups`: The number of groups to split the input into.
      - `channels`: The number of input channels. May be left unspecified (e.g. just
          `None`) if `channelwise_affine=False`.
      - `eps`: Value added to denominator for numerical stability.
      - `channelwise_affine`: Whether the module has learnable affine parameters.
      """
      if (channels is not None) and (channels % groups != 0):
          raise ValueError(
              "The number of groups must divide the number of channels."
          )
      if (channels is None) and channelwise_affine:
          raise ValueError(
              "The number of channels should be specified if "
              "`channelwise_affine=True`"
          )
      self.groups = groups
      self.channels = channels
      self.eps = eps
      self.channelwise_affine = channelwise_affine
      self.weight = jnp.ones(channels) if channelwise_affine else None
      self.bias = jnp.zeros(channels) if channelwise_affine else None
  
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 1 to create a new GroupNorm object from inherited base classes: Module with example: obj = GroupNorm(groups, channels, eps, channelwise_affine)