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","."]
docstr: Normalizes the input using batch statistics.
  
  Args:
    x: the input to be normalized.
    use_running_average: if true, the statistics stored in batch_stats
      will be used instead of computing the batch statistics on the input.
  
  Returns:
    Normalized inputs (the same shape as inputs).
docstr_tok: ["Normalizes","the","input","using","batch","statistics",".","Args",":","x",":","the","input","to","be","normalized",".","use_running_average",":","if","true",",","the","statistics","stored","in","batch_stats","will","be","used","instead","of","computing","the","batch","statistics","on","the","input",".","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","."]
code: def __call__(
      self,
      x,
      use_running_average: tp.Optional[bool] = None,
  ):
      """Normalizes the input using batch statistics.
  
      Args:
        x: the input to be normalized.
        use_running_average: if true, the statistics stored in batch_stats
          will be used instead of computing the batch statistics on the input.
  
      Returns:
        Normalized inputs (the same shape as inputs).
      """
  
      use_running_average = first_from(
          use_running_average,
          self.use_running_average,
          flaglib.flags.get("use_running_average"),
          error_msg="""No `use_running_average` argument was provided to BatchNorm
          as either a __call__ argument, class attribute, or nnx.flag.""",
      )
      feature_axes = _canonicalize_axes(x.ndim, self.axis)
      reduction_axes = tuple(
          i for i in range(x.ndim) if i not in feature_axes
      )
  
      if use_running_average:
          mean, var = self.mean, self.var
      else:
          mean, var = _compute_stats(
              x,
              reduction_axes,
              dtype=self.dtype,
              axis_name=self.axis_name,
              axis_index_groups=self.axis_index_groups,
          )
  
          self.mean = (
              self.momentum * self.mean + (1 - self.momentum) * mean
          )
          self.var = (
              self.momentum * self.var + (1 - self.momentum) * var
          )
  
      return _normalize(
          x,
          mean,
          var,
          self.scale,
          self.bias,
          reduction_axes,
          feature_axes,
          self.dtype,
          self.epsilon,
      )
  
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 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.__call__(x, use_running_average) and returns: _normalize