class_docstr: 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.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> x = jax.random.normal(jax.random.key(0), (5, 6))
    >>> layer = nn.GroupNorm(num_groups=3)
    >>> variables = layer.init(jax.random.key(1), x)
    >>> variables
    {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
    >>> y = layer.apply(variables, x)
  
  Attributes:
    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).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    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.
    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.
class_docstr_tok: ["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",".","Example","usage",":",":",">",">",">","import","flax.linen","as","nn",">",">",">","import","jax",",","jax.numpy","as","jnp",">",">",">","x","=","jax.random.normal","(","jax.random.key","(","0",")",",","(","5",",","6",")",")",">",">",">","layer","=","nn.GroupNorm","(","num_groups=3",")",">",">",">","variables","=","layer.init","(","jax.random.key","(","1",")",",","x",")",">",">",">","variables","{","'params","'",":","{","'scale","'",":","Array","(","[","1.",",","1.",",","1.",",","1.",",","1.",",","1",".","]",",","dtype=float32",")",",","'bias","'",":","Array","(","[","0.",",","0.",",","0.",",","0.",",","0.",",","0",".","]",",","dtype=float32",")","}","}",">",">",">","y","=","layer.apply","(","variables",",","x",")","Attributes",":","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",")",".","param_dtype",":","the","dtype","passed","to","parameter","initializers","(","default",":","float32",")",".","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",".","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","."]
docstr: Applies group normalization to the input (arxiv.org/abs/1803.08494).
  
  Args:
    x: the input of shape ``...C`` where ``C`` 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).
docstr_tok: ["Applies","group","normalization","to","the","input","(","arxiv.org\/abs\/1803.08494",")",".","Args",":","x",":","the","input","of","shape","``","...","C","``","where","``","C","``","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",")","."]
code: def __call__(self, x, mask=None):
      """Applies group normalization to the input (arxiv.org/abs/1803.08494).
  
      Args:
        x: the input of shape ``...C`` where ``C`` 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]
      feature_axis = -1
  
      reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)
  
      if reduction_axes[-1] != (feature_axis % x.ndim):
          raise ValueError(
              "The reduction axes must include the final dimension "
              "as this is assumed to be the feature axis."
          )
  
      if (self.num_groups is None and self.group_size is None) or (
          self.num_groups is not None and self.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."
          )
  
      channels = x.shape[-1]
      if self.group_size is not None:
          if channels % self.group_size != 0:
              raise ValueError(
                  "Number of channels ({}) is not multiple of the "
                  "group size ({}).".format(channels, self.group_size)
              )
          num_groups = channels // self.group_size
      else:
          num_groups = self.num_groups
          assert isinstance(num_groups, int)
  
      if num_groups <= 0 or channels % num_groups != 0:
          raise ValueError(
              "Number of groups ({}) does not divide the number"
              " of channels ({}).".format(num_groups, channels)
          )
  
      group_size = x.shape[-1] // num_groups
      group_shape = x.shape[:-1] + (num_groups, group_size)
  
      if mask is not None:
          mask = mask.reshape(
              mask.shape[:-1] + (num_groups, group_size)
          )
  
      mean, var = _compute_stats(
          x.reshape(group_shape),
          list(reduction_axes[:-1]) + [-1],
          self.dtype,
          self.axis_name,
          self.axis_index_groups,
          use_fast_variance=self.use_fast_variance,
          mask=mask,
      )
      mean = jnp.repeat(mean, group_size, axis=-1)
      var = jnp.repeat(var, group_size, axis=-1)
  
      return _normalize(
          self,
          x,
          mean,
          var,
          reduction_axes[:-1],
          (feature_axis,),
          self.dtype,
          self.param_dtype,
          self.epsilon,
          self.use_bias,
          self.use_scale,
          self.bias_init,
          self.scale_init,
      )
  
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 1 for calling the GroupNorm obj's underlying member method code with example usage: obj.__call__(x, mask) and returns: _normalize