class_docstr: Normalizes inputs along the given axes.
  
  This is a generic implementation of normalization along specific axes of the
  input. :class:`InstanceNorm` is a subclass of this module, it normalizes over
  the spatial dimensions.
  
  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.
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
  
file_path: sonnet/sonnet/src/axis_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 LayerNorm obj's underlying member method code with example usage: obj._initialize(inputs) without return types