class_docstr: Group normalization module.
  
  This applies group normalization to the inputs. 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
      variables.
    - 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.
  
  Attributes:
    scale: If ``create_scale=True``, a trainable :tf:`Variable` holding the
      current scale.
    offset: If ``create_offset=True``, a trainable :tf:`Variable` holding the
      current offset.
class_docstr_tok: ["Group","normalization","module",".","This","applies","group","normalization","to","the","inputs",".","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","variables",".","-","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",".","Attributes",":","scale",":","If","``","create_scale=True","``",",","a","trainable",":","tf",":","`","Variable","`","holding","the","current","scale",".","offset",":","If","``","create_offset=True","``",",","a","trainable",":","tf",":","`","Variable","`","holding","the","current","offset","."]
docstr: Setup of rank specific values.
docstr_tok: ["Setup","of","rank","specific","values","."]
code: def _initialize(self, inputs: tf.Tensor):
      """Setup of rank specific values."""
      self._rank = len(inputs.shape)
  
      # Turns slice into list of axis
      if isinstance(self._axis, slice):
          axes = tuple(range(self._rank))
          self._axis = axes[self._axis]
  
      # Create scale and offset variables
      dtype = inputs.dtype
      if self._channel_index == -1:
          params_shape = [inputs.shape[-1]]
      else:  # self._channel_index == 1
          params_shape = [inputs.shape[1]] + [1] * (self._rank - 2)
  
      if self._create_scale:
          self.scale = tf.Variable(
              self._scale_init(params_shape, dtype), name="scale"
          )
      else:
          self.scale = None
  
      if self._create_offset:
          self.offset = tf.Variable(
              self._offset_init(params_shape, dtype), name="offset"
          )
      else:
          self.offset = None
  
      num_channels = inputs.shape[self._channel_index]
      if num_channels % self._groups != 0:
          raise ValueError(
              "The number of channels must be divisible by the number of groups, "
              "was channels = {}, groups = {}".format(
                  num_channels, self._groups
              )
          )
      if self._channel_index == -1:
          self._inputs_reshape = (
              [-1]
              + list(inputs.shape[1:-1])
              + [self._groups, num_channels // self._groups]
          )
          self._axis = [
              a if a != self._rank - 1 else a + 1 for a in self._axis
          ]
      else:
          self._inputs_reshape = (
              [-1]
              + [self._groups, num_channels // self._groups]
              + list(inputs.shape[2:])
          )
          self._axis = [a if a == 0 else a + 1 for a in self._axis]
      self._outputs_reshape = [-1] + list(inputs.shape[1:])
  
file_path: sonnet/sonnet/src/group_norm.py
imports: import collections.abc
  from typing import Optional
  from sonnet.src import base
  from sonnet.src import initializers
  from sonnet.src import once
  from sonnet.src import types
  from sonnet.src import utils
  import tensorflow
label_desc: Use image node_id 3 for calling the GroupNorm obj's underlying member method code with example usage: obj._initialize(inputs) without return types