class_docstr: Batch normalization module.
  
  This implements normalization across the batch and spatial dimensions.
  It maintains moving averages of the mean and variance which can be
  used to normalize at test time. The constructor is generic and
  requires the user to pass in objects to compute these.
  
  At training time we use the batch statistics for that batch and these are then
  used to update the moving averages.
  
  At test time we can either use the moving averages of the batch statistics
  (``test_local_stats=False``) or we can use the local statistics
  (``test_local_stats=True``).
  
  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``. Note that this module automatically uses the fused batch
  norm op if the data format is ``NHWC``.
  
  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``, a trainable :tf:`Variable` holding the current
      scale after the module is connected for the first time.
    offset: If ``create_offset``, a trainable :tf:`Variable` holding the current
      offset after the module is connected for the first time.
class_docstr_tok: ["Batch","normalization","module",".","This","implements","normalization","across","the","batch","and","spatial","dimensions",".","It","maintains","moving","averages","of","the","mean","and","variance","which","can","be","used","to","normalize","at","test","time",".","The","constructor","is","generic","and","requires","the","user","to","pass","in","objects","to","compute","these",".","At","training","time","we","use","the","batch","statistics","for","that","batch","and","these","are","then","used","to","update","the","moving","averages",".","At","test","time","we","can","either","use","the","moving","averages","of","the","batch","statistics","(","``","test_local_stats=False","``",")","or","we","can","use","the","local","statistics","(","``","test_local_stats=True","``",")",".","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","``",".","Note","that","this","module","automatically","uses","the","fused","batch","norm","op","if","the","data","format","is","``","NHWC","``",".","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","``",",","a","trainable",":","tf",":","`","Variable","`","holding","the","current","scale","after","the","module","is","connected","for","the","first","time",".","offset",":","If","``","create_offset","``",",","a","trainable",":","tf",":","`","Variable","`","holding","the","current","offset","after","the","module","is","connected","for","the","first","time","."]
docstr: Returns normalized inputs.
  
  Args:
    inputs: An n-D tensor of the data_format specified above on which the
      transformation is performed.
    is_training: Whether the module should be connected in training mode,
      meaning the moving averages are updated.
    test_local_stats: Whether local batch statistics should be used when
      ``is_training=False``. If not, moving averages are used. By default
      ``False``.
    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","above","on","which","the","transformation","is","performed",".","is_training",":","Whether","the","module","should","be","connected","in","training","mode",",","meaning","the","moving","averages","are","updated",".","test_local_stats",":","Whether","local","batch","statistics","should","be","used","when","``","is_training=False","``",".","If","not",",","moving","averages","are","used",".","By","default","``","False","``",".","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,
      is_training: types.BoolLike,
      test_local_stats: types.BoolLike = False,
      scale: Optional[tf.Tensor] = None,
      offset: Optional[tf.Tensor] = None,
  ):
      """Returns normalized inputs.
  
      Args:
        inputs: An n-D tensor of the data_format specified above on which the
          transformation is performed.
        is_training: Whether the module should be connected in training mode,
          meaning the moving averages are updated.
        test_local_stats: Whether local batch statistics should be used when
          ``is_training=False``. If not, moving averages are used. By default
          ``False``.
        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.
      """
      use_batch_stats = is_training or test_local_stats
      if self._create_scale:
          if scale is not None:
              raise ValueError(
                  "Cannot pass `scale` at call time if `create_scale=True`."
              )
  
      if self._create_offset:
          if offset is not None:
              raise ValueError(
                  "Cannot pass `offset` at call time if `create_offset=True`."
              )
  
      self._initialize(inputs)
      if scale is None:
          scale = self.scale
      if offset is None:
          offset = self.offset
  
      mean, variance = self._moments(inputs, use_batch_stats)
  
      if self._fused:
          out, mean, variance, _, _ = tf.raw_ops.FusedBatchNormV2(
              x=inputs,
              mean=mean,
              variance=variance,
              scale=scale,
              offset=offset,
              is_training=use_batch_stats,
              epsilon=self._eps,
              data_format=self._fused_data_format,
          )
  
      else:
          out = tf.nn.batch_normalization(
              inputs,
              mean=mean,
              variance=variance,
              scale=scale,
              offset=offset,
              variance_epsilon=self._eps,
          )
  
      if is_training:
          self._update_statistics(mean, variance)
  
      return out
  
file_path: sonnet/sonnet/src/batch_norm.py
imports: from typing import Optional, Tuple
  from sonnet.src import base
  from sonnet.src import initializers
  from sonnet.src import metrics
  from sonnet.src import moving_averages
  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 BaseBatchNorm obj's underlying member method code with example usage: obj.__call__(inputs, is_training, test_local_stats, scale, offset) and returns: out