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:**
  
  - `x`: A JAX array of shape `(input_size, dim_1, ..., dim_N)`.
  - `state`: An [`equinox.nn.State`][] object (which is used to store the
      running statistics).
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
      (Keyword only argument.)
  - `inference`: As per [`equinox.nn.BatchNorm.__init__`][]. If
      `True` or `False` then it will take priority over `self.inference`. If
      `None` then the value from `self.inference` will be used.
  
  **Returns:**
  
  A 2-tuple of:
  
  - A JAX array of shape `(input_size, dim_1, ..., dim_N)`.
  - An updated state object (storing the updated running statistics).
  
  **Raises:**
  
  A `NameError` if no `vmap`s are placed around this operation, or if this vmap
  does not have a matching `axis_name`.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","x","`",":","A","JAX","array","of","shape","`","(","input_size",",","dim_1",",","...",",","dim_N",")","`",".","-","`","state","`",":","An","[","`","equinox.nn.State","`","]","[","]","object","(","which","is","used","to","store","the","running","statistics",")",".","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","(","Keyword","only","argument",".",")","-","`","inference","`",":","As","per","[","`","equinox.nn.BatchNorm.__init__","`","]","[","]",".","If","`","True","`","or","`","False","`","then","it","will","take","priority","over","`","self.inference","`",".","If","`","None","`","then","the","value","from","`","self.inference","`","will","be","used",".","*","*","Returns",":","*","*","A","2-tuple","of",":","-","A","JAX","array","of","shape","`","(","input_size",",","dim_1",",","...",",","dim_N",")","`",".","-","An","updated","state","object","(","storing","the","updated","running","statistics",")",".","*","*","Raises",":","*","*","A","`","NameError","`","if","no","`","vmap","`","s","are","placed","around","this","operation",",","or","if","this","vmap","does","not","have","a","matching","`","axis_name","`","."]
code: def __call__(
      self,
      x: Array,
      state: State,
      *,
      key: Optional[PRNGKeyArray] = None,
      inference: Optional[bool] = None,
  ) -> tuple[Array, State]:
      """**Arguments:**
  
      - `x`: A JAX array of shape `(input_size, dim_1, ..., dim_N)`.
      - `state`: An [`equinox.nn.State`][] object (which is used to store the
          running statistics).
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
          (Keyword only argument.)
      - `inference`: As per [`equinox.nn.BatchNorm.__init__`][]. If
          `True` or `False` then it will take priority over `self.inference`. If
          `None` then the value from `self.inference` will be used.
  
      **Returns:**
  
      A 2-tuple of:
  
      - A JAX array of shape `(input_size, dim_1, ..., dim_N)`.
      - An updated state object (storing the updated running statistics).
  
      **Raises:**
  
      A `NameError` if no `vmap`s are placed around this operation, or if this vmap
      does not have a matching `axis_name`.
      """
  
      if inference is None:
          inference = self.inference
      if inference:
          running_mean, running_var = state.get(self.state_index)
      else:
  
          def _stats(y):
              mean = jnp.mean(y)
              mean = lax.pmean(mean, self.axis_name)
              var = jnp.mean((y - mean) * jnp.conj(y - mean))
              var = lax.pmean(var, self.axis_name)
              var = jnp.maximum(0.0, var)
              return mean, var
  
          first_time = state.get(self.first_time_index)
          state = state.set(self.first_time_index, jnp.array(False))
  
          batch_mean, batch_var = jax.vmap(_stats)(x)
          running_mean, running_var = state.get(self.state_index)
          momentum = self.momentum
          running_mean = (
              1 - momentum
          ) * batch_mean + momentum * running_mean
          running_var = (
              1 - momentum
          ) * batch_var + momentum * running_var
          running_mean = lax.select(
              first_time, batch_mean, running_mean
          )
          running_var = lax.select(first_time, batch_var, running_var)
          state = state.set(
              self.state_index, (running_mean, running_var)
          )
  
      def _norm(y, m, v, w, b):
          out = (y - m) / jnp.sqrt(v + self.eps)
          if self.channelwise_affine:
              out = out * w + b
          return out
  
      out = jax.vmap(_norm)(
          x, running_mean, running_var, self.weight, self.bias
      )
      return out, state
  
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 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.__call__(x, state) and returns: out,  state, out, mean,  var