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: Computes the normalized version of the input.
  
  Args:
    inputs: An array, where the data format is ``[..., C]``.
    is_training: Whether this is during training.
    test_local_stats: Whether local stats are used when is_training=False.
    scale: An array 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: An array 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:
    The array, normalized across all but the last dimension.
docstr_tok: ["Computes","the","normalized","version","of","the","input",".","Args",":","inputs",":","An","array",",","where","the","data","format","is","``","[","...",",","C","]","``",".","is_training",":","Whether","this","is","during","training",".","test_local_stats",":","Whether","local","stats","are","used","when","is_training=False",".","scale",":","An","array","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",":","An","array","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",":","The","array",",","normalized","across","all","but","the","last","dimension","."]
code: def __call__(
      self,
      inputs: jax.Array,
      is_training: bool,
      test_local_stats: bool = False,
      scale: Optional[jax.Array] = None,
      offset: Optional[jax.Array] = None,
  ) -> jax.Array:
      """Computes the normalized version of the input.
  
      Args:
        inputs: An array, where the data format is ``[..., C]``.
        is_training: Whether this is during training.
        test_local_stats: Whether local stats are used when is_training=False.
        scale: An array 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: An array 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:
        The array, normalized across all but the last dimension.
      """
      if self.create_scale and scale is not None:
          raise ValueError(
              "Cannot pass `scale` at call time if `create_scale=True`."
          )
      if self.create_offset and offset is not None:
          raise ValueError(
              "Cannot pass `offset` at call time if `create_offset=True`."
          )
  
      channel_index = self.channel_index
      if channel_index < 0:
          channel_index += inputs.ndim
  
      if self.axis is not None:
          axis = self.axis
      else:
          axis = [i for i in range(inputs.ndim) if i != channel_index]
  
      if is_training or test_local_stats:
          mean = jnp.mean(inputs, axis, keepdims=True)
          mean_of_squares = jnp.mean(
              jnp.square(inputs), axis, keepdims=True
          )
          if self.cross_replica_axis:
              mean = jax.lax.pmean(
                  mean,
                  axis_name=self.cross_replica_axis,
                  axis_index_groups=self.cross_replica_axis_index_groups,
              )
              mean_of_squares = jax.lax.pmean(
                  mean_of_squares,
                  axis_name=self.cross_replica_axis,
                  axis_index_groups=self.cross_replica_axis_index_groups,
              )
          var = mean_of_squares - jnp.square(mean)
      else:
          mean = self.mean_ema.average.astype(inputs.dtype)
          var = self.var_ema.average.astype(inputs.dtype)
  
      if is_training:
          self.mean_ema(mean)
          self.var_ema(var)
  
      w_shape = [
          1 if i in axis else inputs.shape[i]
          for i in range(inputs.ndim)
      ]
      w_dtype = inputs.dtype
  
      if self.create_scale:
          scale = hk.get_parameter(
              "scale", w_shape, w_dtype, self.scale_init
          )
      elif scale is None:
          scale = np.ones([], dtype=w_dtype)
  
      if self.create_offset:
          offset = hk.get_parameter(
              "offset", w_shape, w_dtype, self.offset_init
          )
      elif offset is None:
          offset = np.zeros([], dtype=w_dtype)
  
      eps = jax.lax.convert_element_type(self.eps, var.dtype)
      inv = scale * jax.lax.rsqrt(var + eps)
      return (inputs - mean) * inv + offset
  
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 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.__call__(inputs, is_training, test_local_stats, scale, offset) and returns: unknown