class_docstr: Group normalization module.
  
  This applies group normalization to the x. This involves splitting the
  channels into groups before calculating the mean and variance. The default
  behaviour is to compute the mean and variance over the spatial dimensions and
  the grouped channels. The mean and variance will never be computed over the
  created groups axis.
  
  It transforms the input ``x`` into:
  
  .. math::
  
     \d{outputs} = \d{scale} \dfrac{x - \mu}{\sigma + \epsilon} + \d{offset}
  
  Where :math:`\mu` and :math:`\sigma` are respectively the mean and standard
  deviation of ``x``.
  
  There are many different variations for how users want to manage scale and
  offset if they require them at all. These are:
  
    - No ``scale``/``offset`` in which case ``create_*`` should be set to
      ``False`` and ``scale``/``offset`` aren't passed when the module is
      called.
    - Trainable ``scale``/``offset`` in which case create_* should be set to
      ``True`` and again ``scale``/``offset`` aren't passed when the module is
      called. In this case this module creates and owns the scale/offset
      parameters.
    - Externally generated ``scale``/``offset``, such as for conditional
      normalization, in which case ``create_*`` should be set to ``False`` and
      then the values fed in at call time.
class_docstr_tok: ["Group","normalization","module",".","This","applies","group","normalization","to","the","x",".","This","involves","splitting","the","channels","into","groups","before","calculating","the","mean","and","variance",".","The","default","behaviour","is","to","compute","the","mean","and","variance","over","the","spatial","dimensions","and","the","grouped","channels",".","The","mean","and","variance","will","never","be","computed","over","the","created","groups","axis",".","It","transforms","the","input","``","x","``","into",":","..","math",":",":","\\d","{","outputs","}","=","\\d","{","scale","}","\\dfrac","{","x","-","\\mu","}","{","\\sigma","+","\\epsilon","}","+","\\d","{","offset","}","Where",":","math",":","`","\\mu","`","and",":","math",":","`","\\sigma","`","are","respectively","the","mean","and","standard","deviation","of","``","x","``",".","There","are","many","different","variations","for","how","users","want","to","manage","scale","and","offset","if","they","require","them","at","all",".","These","are",":","-","No","``","scale","``","\/","``","offset","``","in","which","case","``","create_","*","``","should","be","set","to","``","False","``","and","``","scale","``","\/","``","offset","``","are","n't","passed","when","the","module","is","called",".","-","Trainable","``","scale","``","\/","``","offset","``","in","which","case","create_","*","should","be","set","to","``","True","``","and","again","``","scale","``","\/","``","offset","``","are","n't","passed","when","the","module","is","called",".","In","this","case","this","module","creates","and","owns","the","scale\/offset","parameters",".","-","Externally","generated","``","scale","``","\/","``","offset","``",",","such","as","for","conditional","normalization",",","in","which","case","``","create_","*","``","should","be","set","to","``","False","``","and","then","the","values","fed","in","at","call","time","."]
docstr: Returns normalized inputs.
  
  Args:
    x: An n-D tensor of the ``data_format`` specified in the constructor
      on which the transformation is performed.
    scale: A tensor up to n-D. The shape of this tensor must be broadcastable
      to the shape of ``x``. This is the scale applied to the normalized
      x. This cannot be passed in if the module was constructed with
      ``create_scale=True``.
    offset: A tensor up to n-D. The shape of this tensor must be broadcastable
      to the shape of ``x``. This is the offset applied to the normalized
      ``x``. This cannot be passed in if the module was constructed with
      ``create_offset=True``.
  
  Returns:
    An n-d tensor of the same shape as x that has been normalized.
docstr_tok: ["Returns","normalized","inputs",".","Args",":","x",":","An","n-D","tensor","of","the","``","data_format","``","specified","in","the","constructor","on","which","the","transformation","is","performed",".","scale",":","A","tensor","up","to","n-D",".","The","shape","of","this","tensor","must","be","broadcastable","to","the","shape","of","``","x","``",".","This","is","the","scale","applied","to","the","normalized","x",".","This","can","not","be","passed","in","if","the","module","was","constructed","with","``","create_scale=True","``",".","offset",":","A","tensor","up","to","n-D",".","The","shape","of","this","tensor","must","be","broadcastable","to","the","shape","of","``","x","``",".","This","is","the","offset","applied","to","the","normalized","``","x","``",".","This","can","not","be","passed","in","if","the","module","was","constructed","with","``","create_offset=True","``",".","Returns",":","An","n-d","tensor","of","the","same","shape","as","x","that","has","been","normalized","."]
code: def __call__(
      self,
      x: jax.Array,
      scale: Optional[jax.Array] = None,
      offset: Optional[jax.Array] = None,
  ) -> jax.Array:
      """Returns normalized inputs.
  
      Args:
        x: An n-D tensor of the ``data_format`` specified in the constructor
          on which the transformation is performed.
        scale: A tensor up to n-D. The shape of this tensor must be broadcastable
          to the shape of ``x``. This is the scale applied to the normalized
          x. This cannot be passed in if the module was constructed with
          ``create_scale=True``.
        offset: A tensor up to n-D. The shape of this tensor must be broadcastable
          to the shape of ``x``. This is the offset applied to the normalized
          ``x``. This cannot be passed in if the module was constructed with
          ``create_offset=True``.
  
      Returns:
        An n-d tensor of the same shape as x that has been normalized.
      """
      if self.rank is not None and x.ndim != self.rank:
          raise ValueError(
              "The rank of the inputs cannot change between calls, the"
              f" original call was rank={self.rank} but this call was "
              f"rank={x.ndim}."
          )
  
      if self.create_scale and scale is not None:
          raise ValueError(
              "Cannot pass `scale` at call time if `create_scale=True`."
          )
  
      if self.create_offset and offset is not None:
          raise ValueError(
              "Cannot pass `offset` at call time if `create_offset=True`."
          )
  
      channels = x.shape[self.channel_index]
      if channels % self.groups != 0:
          raise ValueError(
              "The number of channels must be divisible by the number of groups, "
              f"was channels={channels}, groups={self.groups}"
          )
  
      if self.rank is None:
          self._initialize(x, channels)
  
      dtype = x.dtype
      if self.channel_index == -1:
          params_shape = (x.shape[-1],)
      else:
          assert self.channel_index == 1
          params_shape = (x.shape[1],) + (1,) * (self.rank - 2)
  
      if self.create_scale:
          scale = hk.get_parameter(
              "scale", params_shape, dtype, self.scale_init
          )
  
      if self.create_offset:
          offset = hk.get_parameter(
              "offset", params_shape, dtype, self.offset_init
          )
  
      x = x.reshape(self.group_shape)
      mean = jnp.mean(x, self.axis, keepdims=True)
      # TODO(tycai): Consider faster but less precise variance formulation.
      var = jnp.var(x, self.axis, keepdims=True)
      x = (x - mean) * jax.lax.rsqrt(var + self.eps)
      x = x.reshape(self.first_input_shape)
  
      if scale is not None:
          scale = jax.lax.broadcast_to_rank(scale, x.ndim)
          x = x * scale
  
      if offset is not None:
          offset = jax.lax.broadcast_to_rank(offset, x.ndim)
          x = x + offset
  
      return x
  
file_path: dm-haiku/haiku/_src/group_norm.py
imports: import collections
  from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the GroupNorm obj's underlying member method code with example usage: obj.__call__(x, scale, offset) and returns: x