class_docstr: BatchNorm Module.
  
  Attributes:
    use_running_average: if True, the statistics stored in batch_stats
      will be used instead of computing the batch statistics on the input.
    axis: the feature or non-batch axis of the input.
    momentum: decay rate for the exponential moving average of
      the batch statistics.
    epsilon: a small float added to variance to avoid dividing by zero.
    dtype: the dtype of the result (default: infer from input and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    use_bias:  if True, bias (beta) is added.
    use_scale: if True, multiply by scale (gamma).
      When the next layer is linear (also e.g. nn.relu), this can be disabled
      since the scaling will be done by the next layer.
    bias_init: initializer for bias, by default, zero.
    scale_init: initializer for scale, by default, one.
    axis_name: the axis name used to combine batch statistics from multiple
      devices. See `jax.pmap` for a description of axis names (default: None).
    axis_index_groups: groups of axis indices within that named axis
      representing subsets of devices to reduce over (default: None). For
      example, `[[0, 1], [2, 3]]` would independently batch-normalize over
      the examples on the first two and last two devices. See `jax.lax.psum`
      for more details.
class_docstr_tok: ["BatchNorm","Module",".","Attributes",":","use_running_average",":","if","True",",","the","statistics","stored","in","batch_stats","will","be","used","instead","of","computing","the","batch","statistics","on","the","input",".","axis",":","the","feature","or","non-batch","axis","of","the","input",".","momentum",":","decay","rate","for","the","exponential","moving","average","of","the","batch","statistics",".","epsilon",":","a","small","float","added","to","variance","to","avoid","dividing","by","zero",".","dtype",":","the","dtype","of","the","result","(","default",":","infer","from","input","and","params",")",".","param_dtype",":","the","dtype","passed","to","parameter","initializers","(","default",":","float32",")",".","use_bias",":","if","True",",","bias","(","beta",")","is","added",".","use_scale",":","if","True",",","multiply","by","scale","(","gamma",")",".","When","the","next","layer","is","linear","(","also","e.g",".","nn.relu",")",",","this","can","be","disabled","since","the","scaling","will","be","done","by","the","next","layer",".","bias_init",":","initializer","for","bias",",","by","default",",","zero",".","scale_init",":","initializer","for","scale",",","by","default",",","one",".","axis_name",":","the","axis","name","used","to","combine","batch","statistics","from","multiple","devices",".","See","`","jax.pmap","`","for","a","description","of","axis","names","(","default",":","None",")",".","axis_index_groups",":","groups","of","axis","indices","within","that","named","axis","representing","subsets","of","devices","to","reduce","over","(","default",":","None",")",".","For","example",",","`","[","[","0",",","1","]",",","[","2",",","3","]","]","`","would","independently","batch-normalize","over","the","examples","on","the","first","two","and","last","two","devices",".","See","`","jax.lax.psum","`","for","more","details","."]
code: def __init__(
      self,
      num_features: int,
      *,
      use_running_average: tp.Optional[bool] = None,
      axis: int = -1,
      momentum: float = 0.99,
      epsilon: float = 1e-5,
      dtype: tp.Optional[Dtype] = None,
      param_dtype: Dtype = jnp.float32,
      use_bias: bool = True,
      use_scale: bool = True,
      bias_init: tp.Callable[
          [KeyArray, Shape, Dtype], Array
      ] = initializers.zeros_init(),
      scale_init: tp.Callable[
          [KeyArray, Shape, Dtype], Array
      ] = initializers.ones_init(),
      axis_name: tp.Optional[str] = None,
      axis_index_groups: tp.Any = None,
      rngs: rnglib.Rngs,
  ):
      feature_shape = (num_features,)
      self.mean = nnx.BatchStat(jnp.zeros(feature_shape, jnp.float32))
      self.var = nnx.BatchStat(jnp.ones(feature_shape, jnp.float32))
  
      if use_scale:
          key = rngs.params()
          self.scale = nnx.Param(
              scale_init(key, feature_shape, param_dtype)
          )
      else:
          self.scale = nnx.Param(None)
  
      if use_bias:
          key = rngs.params()
          self.bias = nnx.Param(
              bias_init(key, feature_shape, param_dtype)
          )
      else:
          self.bias = nnx.Param(None)
  
      self.num_features = num_features
      self.use_running_average = use_running_average
      self.axis = axis
      self.momentum = momentum
      self.epsilon = epsilon
      self.dtype = dtype
      self.param_dtype = param_dtype
      self.use_bias = use_bias
      self.use_scale = use_scale
      self.bias_init = bias_init
      self.scale_init = scale_init
      self.axis_name = axis_name
      self.axis_index_groups = axis_index_groups
  
file_path: flax/flax/experimental/nnx/nnx/nn/normalization.py
imports: import typing
  import jax
  import jax.numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib, rnglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 1 to create a new BatchNorm object from inherited base classes: Module with example: obj = BatchNorm(num_features)