class_docstr: Layer normalization (https://arxiv.org/abs/1607.06450) with named axes for :meth:`jax.experimental.pjit.pjit`.
  
  .. warning:: This class is hightly EXPERIMENTAL and the API is likely to
      change. For regular (non-pjit) use, please use
      :class:`flax.linen.normalization.LayerNorm`.
  
  Operates on the last axis of the input data.
  
  It 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 computation (default: float32).
    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.
docstr: Applies layer normalization on the input.
  
  Args:
    x: the inputs
  
  Returns:
    Normalized inputs (the same shape as inputs).
docstr_tok: ["Applies","layer","normalization","on","the","input",".","Args",":","x",":","the","inputs","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","."]
code: def __call__(self, x):
      """Applies layer normalization on the input.
  
      Args:
        x: the inputs
  
      Returns:
        Normalized inputs (the same shape as inputs).
      """
      reduction_axes = (-1,)
      feature_axes = (-1,)
  
      mean, var = _compute_stats(x, reduction_axes)
  
      return _normalize(
          self,
          x,
          mean,
          var,
          reduction_axes,
          feature_axes,
          self.dtype,
          self.param_dtype,
          self.epsilon,
          self.use_bias,
          self.use_scale,
          self.bias_init,
          self.scale_init,
      )
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import jax.numpy
  from jax import lax
  from flax import linen
  from flax.linen import initializers
  from flax.linen.linear import DotGeneralT, PrecisionLike
  from flax.linen.partitioning import param_with_axes, with_sharding_constraint
label_desc: Use image node_id 1 for calling the LayerNorm obj's underlying member method code with example usage: obj.__call__(x) and returns: _normalize