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.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> x = jax.random.normal(jax.random.key(0), (5, 6))
    >>> layer = nn.LayerNorm()
    >>> variables = layer.init(jax.random.key(1), x)
    >>> variables
    {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
    >>> y = layer.apply(variables, x)
  
  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 or shard
      map. For SPMD jit, you do not need to manually synchronize. Just make sure
      that the axes are correctly annotated and XLA:SPMD will insert the
      necessary collectives.
    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.
    use_fast_variance: If true, use a faster, but less numerically stable,
      calculation for the variance.
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",".","Example","usage",":",":",">",">",">","import","flax.linen","as","nn",">",">",">","import","jax",",","jax.numpy","as","jnp",">",">",">","x","=","jax.random.normal","(","jax.random.key","(","0",")",",","(","5",",","6",")",")",">",">",">","layer","=","nn.LayerNorm","(",")",">",">",">","variables","=","layer.init","(","jax.random.key","(","1",")",",","x",")",">",">",">","variables","{","'params","'",":","{","'scale","'",":","Array","(","[","1.",",","1.",",","1.",",","1.",",","1.",",","1",".","]",",","dtype=float32",")",",","'bias","'",":","Array","(","[","0.",",","0.",",","0.",",","0.",",","0.",",","0",".","]",",","dtype=float32",")","}","}",">",">",">","y","=","layer.apply","(","variables",",","x",")","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","or","shard","map",".","For","SPMD","jit",",","you","do","not","need","to","manually","synchronize",".","Just","make","sure","that","the","axes","are","correctly","annotated","and","XLA",":","SPMD","will","insert","the","necessary","collectives",".","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",".","use_fast_variance",":","If","true",",","use","a","faster",",","but","less","numerically","stable",",","calculation","for","the","variance","."]
docstr: Applies layer normalization on the input.
  
  Args:
    x: the inputs
    mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating
      the positions for which the mean and variance should be computed.
  
  Returns:
    Normalized inputs (the same shape as inputs).
docstr_tok: ["Applies","layer","normalization","on","the","input",".","Args",":","x",":","the","inputs","mask",":","Binary","array","of","shape","broadcastable","to","``","inputs","``","tensor",",","indicating","the","positions","for","which","the","mean","and","variance","should","be","computed",".","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","."]
code: def __call__(self, x, mask=None):
      """Applies layer normalization on the input.
  
      Args:
        x: the inputs
        mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating
          the positions for which the mean and variance should be computed.
  
      Returns:
        Normalized inputs (the same shape as inputs).
      """
      mean, var = _compute_stats(
          x,
          self.reduction_axes,
          self.dtype,
          self.axis_name,
          self.axis_index_groups,
          use_fast_variance=self.use_fast_variance,
          mask=mask,
      )
  
      return _normalize(
          self,
          x,
          mean,
          var,
          self.reduction_axes,
          self.feature_axes,
          self.dtype,
          self.param_dtype,
          self.epsilon,
          self.use_bias,
          self.use_scale,
          self.bias_init,
          self.scale_init,
      )
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  from jax import lax
  from jax.nn import initializers
  from flax.linen import dtypes, module, transforms
label_desc: Use image node_id 1 for calling the LayerNorm obj's underlying member method code with example usage: obj.__call__(x, mask) and returns: _normalize