class_docstr: Normalizes inputs to maintain a mean of ~0 and stddev of ~1.
  
  See: https://arxiv.org/abs/1502.03167.
  
  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.
  
  NOTE: ``jax.vmap(hk.transform(BatchNorm))`` will update summary statistics and
  normalize values on a per-batch basis; we currently do *not* support
  normalizing across a batch axis introduced by vmap.
class_docstr_tok: ["Normalizes","inputs","to","maintain","a","mean","of","~0","and","stddev","of","~1",".","See",":","https",":","\/\/arxiv.org\/abs\/1502.03167",".","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",".","NOTE",":","``","jax.vmap","(","hk.transform","(","BatchNorm",")",")","``","will","update","summary","statistics","and","normalize","values","on","a","per-batch","basis",";","we","currently","do","*","not","*","support","normalizing","across","a","batch","axis","introduced","by","vmap","."]
docstr: Constructs a BatchNorm module.
  
  Args:
    create_scale: Whether to include a trainable scaling factor.
    create_offset: Whether to include a trainable offset.
    decay_rate: Decay rate for EMA.
    eps: Small epsilon to avoid division by zero variance. Defaults ``1e-5``,
      as in the paper and Sonnet.
    scale_init: Optional initializer for gain (aka scale). Can only be set
      if ``create_scale=True``. By default, ``1``.
    offset_init: Optional initializer for bias (aka offset). Can only be set
      if ``create_offset=True``. By default, ``0``.
    axis: Which axes to reduce over. The default (``None``) signifies that all
      but the channel axis should be normalized. Otherwise this is a list of
      axis indices which will have normalization statistics calculated.
    cross_replica_axis: If not ``None``, it should be a string (or sequence of
      strings) representing the axis name(s) over which this module is being
      run within a jax map (e.g. ``jax.pmap`` or ``jax.vmap``). Supplying this
      argument means that batch statistics are calculated across all replicas
      on the named axes.
    cross_replica_axis_index_groups: Specifies how devices are grouped. Valid
      only within ``jax.pmap`` collectives.
    data_format: The data format of the input. Can be either
      ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
      default it is ``channels_last``. See :func:`get_channel_index`.
    name: The module name.
docstr_tok: ["Constructs","a","BatchNorm","module",".","Args",":","create_scale",":","Whether","to","include","a","trainable","scaling","factor",".","create_offset",":","Whether","to","include","a","trainable","offset",".","decay_rate",":","Decay","rate","for","EMA",".","eps",":","Small","epsilon","to","avoid","division","by","zero","variance",".","Defaults","``","1e-5","``",",","as","in","the","paper","and","Sonnet",".","scale_init",":","Optional","initializer","for","gain","(","aka","scale",")",".","Can","only","be","set","if","``","create_scale=True","``",".","By","default",",","``","1","``",".","offset_init",":","Optional","initializer","for","bias","(","aka","offset",")",".","Can","only","be","set","if","``","create_offset=True","``",".","By","default",",","``","0","``",".","axis",":","Which","axes","to","reduce","over",".","The","default","(","``","None","``",")","signifies","that","all","but","the","channel","axis","should","be","normalized",".","Otherwise","this","is","a","list","of","axis","indices","which","will","have","normalization","statistics","calculated",".","cross_replica_axis",":","If","not","``","None","``",",","it","should","be","a","string","(","or","sequence","of","strings",")","representing","the","axis","name","(","s",")","over","which","this","module","is","being","run","within","a","jax","map","(","e.g",".","``","jax.pmap","``","or","``","jax.vmap","``",")",".","Supplying","this","argument","means","that","batch","statistics","are","calculated","across","all","replicas","on","the","named","axes",".","cross_replica_axis_index_groups",":","Specifies","how","devices","are","grouped",".","Valid","only","within","``","jax.pmap","``","collectives",".","data_format",":","The","data","format","of","the","input",".","Can","be","either","``","channels_first","``",",","``","channels_last","``",",","``","N","...","C","``","or","``","NC","...","``",".","By","default","it","is","``","channels_last","``",".","See",":","func",":","`","get_channel_index","`",".","name",":","The","module","name","."]
code: def __init__(
      self,
      create_scale: bool,
      create_offset: bool,
      decay_rate: float,
      eps: float = 1e-5,
      scale_init: Optional[hk.initializers.Initializer] = None,
      offset_init: Optional[hk.initializers.Initializer] = None,
      axis: Optional[Sequence[int]] = None,
      cross_replica_axis: Optional[Union[str, Sequence[str]]] = None,
      cross_replica_axis_index_groups: Optional[
          Sequence[Sequence[int]]
      ] = None,
      data_format: str = "channels_last",
      name: Optional[str] = None,
  ):
      """Constructs a BatchNorm module.
  
      Args:
        create_scale: Whether to include a trainable scaling factor.
        create_offset: Whether to include a trainable offset.
        decay_rate: Decay rate for EMA.
        eps: Small epsilon to avoid division by zero variance. Defaults ``1e-5``,
          as in the paper and Sonnet.
        scale_init: Optional initializer for gain (aka scale). Can only be set
          if ``create_scale=True``. By default, ``1``.
        offset_init: Optional initializer for bias (aka offset). Can only be set
          if ``create_offset=True``. By default, ``0``.
        axis: Which axes to reduce over. The default (``None``) signifies that all
          but the channel axis should be normalized. Otherwise this is a list of
          axis indices which will have normalization statistics calculated.
        cross_replica_axis: If not ``None``, it should be a string (or sequence of
          strings) representing the axis name(s) over which this module is being
          run within a jax map (e.g. ``jax.pmap`` or ``jax.vmap``). Supplying this
          argument means that batch statistics are calculated across all replicas
          on the named axes.
        cross_replica_axis_index_groups: Specifies how devices are grouped. Valid
          only within ``jax.pmap`` collectives.
        data_format: The data format of the input. Can be either
          ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
          default it is ``channels_last``. See :func:`get_channel_index`.
        name: The module name.
      """
      super().__init__(name=name)
      if not create_scale and scale_init is not None:
          raise ValueError(
              "Cannot set `scale_init` if `create_scale=False`"
          )
      if not create_offset and offset_init is not None:
          raise ValueError(
              "Cannot set `offset_init` if `create_offset=False`"
          )
      if (
          cross_replica_axis is None
          and cross_replica_axis_index_groups is not None
      ):
          raise ValueError(
              "`cross_replica_axis` name must be specified"
              "if `cross_replica_axis_index_groups` are used."
          )
  
      self.create_scale = create_scale
      self.create_offset = create_offset
      self.eps = eps
      self.scale_init = scale_init or jnp.ones
      self.offset_init = offset_init or jnp.zeros
      self.axis = axis
      self.cross_replica_axis = cross_replica_axis
      self.cross_replica_axis_index_groups = (
          cross_replica_axis_index_groups
      )
      self.channel_index = hk.get_channel_index(data_format)
      self.mean_ema = hk.ExponentialMovingAverage(
          decay_rate, name="mean_ema"
      )
      self.var_ema = hk.ExponentialMovingAverage(
          decay_rate, name="var_ema"
      )
  
file_path: dm-haiku/haiku/_src/batch_norm.py
imports: 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 moving_averages
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 to create a new BatchNorm object from inherited base classes: hk with example: obj = BatchNorm(create_scale, create_offset, decay_rate, eps, scale_init, offset_init, axis, cross_replica_axis, cross_replica_axis_index_groups, data_format, name)