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: Returns normalized inputs.
  
  Args:
    inputs: 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 ``inputs``. This is the scale applied to the normalized
      inputs. 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 ``inputs``. This is the offset applied to the normalized
      ``inputs``. This cannot be passed in if the module was constructed with
      ``create_offset=True``.
  
  Returns:
    An n-d tensor of the same shape as inputs that has been normalized.
docstr_tok: ["Returns","normalized","inputs",".","Args",":","inputs",":","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","``","inputs","``",".","This","is","the","scale","applied","to","the","normalized","inputs",".","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","``","inputs","``",".","This","is","the","offset","applied","to","the","normalized","``","inputs","``",".","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","inputs","that","has","been","normalized","."]
code: def __call__(
      self,
      inputs: tf.Tensor,
      scale: Optional[tf.Tensor] = None,
      offset: Optional[tf.Tensor] = None,
  ) -> tf.Tensor:
      """Returns normalized inputs.
  
      Args:
        inputs: 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 ``inputs``. This is the scale applied to the normalized
          inputs. 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 ``inputs``. This is the offset applied to the normalized
          ``inputs``. This cannot be passed in if the module was constructed with
          ``create_offset=True``.
  
      Returns:
        An n-d tensor of the same shape as inputs that has been normalized.
      """
      self._initialize(inputs)
      if self._create_scale:
          if scale is not None:
              raise ValueError(
                  "Cannot pass `scale` at call time if `create_scale=True`."
              )
          scale = self.scale
  
      if self._create_offset:
          if offset is not None:
              raise ValueError(
                  "Cannot pass `offset` at call time if `create_offset=True`."
              )
          offset = self.offset
  
      if len(inputs.shape) != self._rank:
          raise ValueError(
              "The rank of the inputs cannot change between calls, the"
              " original call was rank={} but this call was rank={}.".format(
                  self._rank, len(inputs.shape)
              )
          )
  
      mean, var = tf.nn.moments(inputs, self._axis, keepdims=True)
  
      normalized = tf.nn.batch_normalization(
          inputs,
          mean=mean,
          variance=var,
          scale=scale,
          offset=offset,
          variance_epsilon=self._eps,
      )
      return normalized
  
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 2 for calling the LayerNorm obj's underlying member method code with example usage: obj.__call__(inputs, scale, offset) and returns: normalized