class_docstr: Layer normalization (https://arxiv.org/abs/1607.06450).
  
  LayerNorm normalizes the activations of the layer for each given example in a
  batch independently, rather than across a batch like Batch Normalization.
  i.e. applies a transformation that maintains the mean activation within
  each example close to 0 and the activation standard deviation close to 1.
  
  Attributes:
      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.
      reduction_axes: Axes for computing normalization statistics.
      feature_axes: Feature axes for learned bias and scaling.
      axis_name: the axis name used to combine batch statistics from multiple
          devices. See `jax.pmap` for a description of axis names (default: None).
          This is only needed if the model is subdivided across devices, i.e. the
          array being normalized is sharded across devices within a pmap.
      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: ["Layer","normalization","(","https",":","\/\/arxiv.org\/abs\/1607.06450",")",".","LayerNorm","normalizes","the","activations","of","the","layer","for","each","given","example","in","a","batch","independently",",","rather","than","across","a","batch","like","Batch","Normalization",".","i.e",".","applies","a","transformation","that","maintains","the","mean","activation","within","each","example","close","to","0","and","the","activation","standard","deviation","close","to","1",".","Attributes",":","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",".","reduction_axes",":","Axes","for","computing","normalization","statistics",".","feature_axes",":","Feature","axes","for","learned","bias","and","scaling",".","axis_name",":","the","axis","name","used","to","combine","batch","statistics","from","multiple","devices",".","See","`","jax.pmap","`","for","a","description","of","axis","names","(","default",":","None",")",".","This","is","only","needed","if","the","model","is","subdivided","across","devices",",","i.e",".","the","array","being","normalized","is","sharded","across","devices","within","a","pmap",".","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,
      *,
      epsilon: float = 1e-6,
      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(),
      reduction_axes: Axes = -1,
      feature_axes: Axes = -1,
      axis_name: tp.Optional[str] = None,
      axis_index_groups: tp.Any = None,
      rngs: rnglib.Rngs,
  ):
      feature_shape = (num_features,)
  
      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.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.reduction_axes = reduction_axes
      self.feature_axes = feature_axes
      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 LayerNorm object from inherited base classes: Module with example: obj = LayerNorm(num_features)