class_docstr: Computes a mean and standard deviation over the batch and spatial
  dimensions of an array, and uses these to normalise the whole array. Optionally
  applies a channelwise affine transformation afterwards.
  
  Given an input array $x = [x_1, ... x_C]$ with $C$ channels, this layer computes
  
  $$\frac{x_i - \mathbb{E}[x_i]}{\sqrt{\text{Var}[x_i] + \varepsilon}} * \gamma_i + \beta_i$$
  
  for all $i$. Here $*$ denotes elementwise multiplication and $\gamma$, $\beta$ have
  shape $(C,)$ if `channelwise_affine=True` and $\gamma = 1$, $\beta = 0$ if
  `channelwise_affine=False`. Expectations are computed over all spatial dimensions
  *and* over the batch dimension, and updated batch-by-batch according to `momentum`.
  
  !!! example
  
      See [this example](../../examples/stateful.ipynb) for example usage.
  
  !!! warning
  
      This layer must be used inside of a `vmap` or `pmap` with a matching
      `axis_name`. (Not doing so will raise a `NameError`.)
  
  Note that this layer behaves differently during training and inference. During
  training then statistics are computed using the input data, and the running
  statistics updated. During inference then just the running statistics are used.
  Whether the model is in training or inference mode should be toggled using
  [`equinox.nn.inference_mode`][].
docstr: **Arguments:**
  
  - `input_size`: The number of channels in the input array.
  - `axis_name`: The name of the batch axis to compute statistics over, as passed
      to `axis_name` in `jax.vmap` or `jax.pmap`. Can also be a sequence (e.g. a
      tuple or a list) of names, to compute statistics over multiple named axes.
  - `eps`: Value added to the denominator for numerical stability.
  - `channelwise_affine`: Whether the module has learnable channel-wise affine
      parameters.
  - `momentum`: The rate at which to update the running statistics. Should be a
      value between 0 and 1 exclusive.
  - `inference`: If `False` then the batch means and variances will be calculated
      and used to update the running statistics. If `True` then the running
      statistics are directly used for normalisation. This may be toggled with
      [`equinox.nn.inference_mode`][] or overridden during
      [`equinox.nn.BatchNorm.__call__`][].
  - `dtype`: The dtype to use for the running statistics. Defaults to either
      `jax.numpy.float32` or `jax.numpy.float64` depending on whether JAX is in
      64-bit mode.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","input_size","`",":","The","number","of","channels","in","the","input","array",".","-","`","axis_name","`",":","The","name","of","the","batch","axis","to","compute","statistics","over",",","as","passed","to","`","axis_name","`","in","`","jax.vmap","`","or","`","jax.pmap","`",".","Can","also","be","a","sequence","(","e.g",".","a","tuple","or","a","list",")","of","names",",","to","compute","statistics","over","multiple","named","axes",".","-","`","eps","`",":","Value","added","to","the","denominator","for","numerical","stability",".","-","`","channelwise_affine","`",":","Whether","the","module","has","learnable","channel-wise","affine","parameters",".","-","`","momentum","`",":","The","rate","at","which","to","update","the","running","statistics",".","Should","be","a","value","between","0","and","1","exclusive",".","-","`","inference","`",":","If","`","False","`","then","the","batch","means","and","variances","will","be","calculated","and","used","to","update","the","running","statistics",".","If","`","True","`","then","the","running","statistics","are","directly","used","for","normalisation",".","This","may","be","toggled","with","[","`","equinox.nn.inference_mode","`","]","[","]","or","overridden","during","[","`","equinox.nn.BatchNorm.__call__","`","]","[","]",".","-","`","dtype","`",":","The","dtype","to","use","for","the","running","statistics",".","Defaults","to","either","`","jax.numpy.float32","`","or","`","jax.numpy.float64","`","depending","on","whether","JAX","is","in","64-bit","mode","."]
code: def __init__(
      self,
      input_size: int,
      axis_name: Union[Hashable, Sequence[Hashable]],
      eps: float = 1e-5,
      channelwise_affine: bool = True,
      momentum: float = 0.99,
      inference: bool = False,
      dtype=None,
  ):
      """**Arguments:**
  
      - `input_size`: The number of channels in the input array.
      - `axis_name`: The name of the batch axis to compute statistics over, as passed
          to `axis_name` in `jax.vmap` or `jax.pmap`. Can also be a sequence (e.g. a
          tuple or a list) of names, to compute statistics over multiple named axes.
      - `eps`: Value added to the denominator for numerical stability.
      - `channelwise_affine`: Whether the module has learnable channel-wise affine
          parameters.
      - `momentum`: The rate at which to update the running statistics. Should be a
          value between 0 and 1 exclusive.
      - `inference`: If `False` then the batch means and variances will be calculated
          and used to update the running statistics. If `True` then the running
          statistics are directly used for normalisation. This may be toggled with
          [`equinox.nn.inference_mode`][] or overridden during
          [`equinox.nn.BatchNorm.__call__`][].
      - `dtype`: The dtype to use for the running statistics. Defaults to either
          `jax.numpy.float32` or `jax.numpy.float64` depending on whether JAX is in
          64-bit mode.
      """
  
      if channelwise_affine:
          self.weight = jnp.ones((input_size,))
          self.bias = jnp.zeros((input_size,))
      else:
          self.weight = None
          self.bias = None
      self.first_time_index = StateIndex(jnp.array(True))
      if dtype is None:
          dtype = default_floating_dtype()
      init_buffers = (
          jnp.empty((input_size,), dtype=dtype),
          jnp.empty((input_size,), dtype=dtype),
      )
      self.state_index = StateIndex(init_buffers)
      self.inference = inference
      self.axis_name = axis_name
      self.input_size = input_size
      self.eps = eps
      self.channelwise_affine = channelwise_affine
      self.momentum = momentum
  
file_path: equinox/equinox/nn/_batch_norm.py
imports: from collections.abc import Hashable, Sequence
  from typing import Optional, Union
  import jax
  import jax.lax
  import jax.numpy
  from jaxtyping import Array, Bool, Float, PRNGKeyArray
  from .._misc import default_floating_dtype
  from .._module import field
  from ._sequential import StatefulLayer
  from ._stateful import State, StateIndex
label_desc: Use image node_id 1 to create a new BatchNorm object from inherited base classes: StatefulLayer with example: obj = BatchNorm(input_size, axis_name, eps, channelwise_affine, momentum, inference, dtype)