class_docstr: Computes a mean and standard deviation over the whole input array, and uses these
  to normalise the whole array. Optionally applies an elementwise affine
  transformation afterwards.
  
  Given an input array $x$, this layer computes
  
  $$\frac{x - \mathbb{E}[x]}{\sqrt{\text{Var}[x] + \varepsilon}} * \gamma + \beta$$
  
  where $\gamma$, $\beta$ have the same shape as $x$ if `elementwise_affine=True`,
  and $\gamma = 1$, $\beta = 0$ if `elementwise_affine=False`.
  
  ??? cite
      [Layer Normalization](https://arxiv.org/abs/1607.06450)
  
      ```bibtex
      @article{ba2016layer,
          author={Jimmy Lei Ba, Jamie Ryan Kriso, Geoffrey E. Hinton},
          title={Layer Normalization},
          year={2016},
          journal={arXiv:1607.06450},
      }
      ```
  
  !!! faq "FAQ"
  
      If you need to normalise over only some input dimensions, then this can be
      achieved by vmap'ing. For example the following will compute statistics over
      every dimension *except* the first:
      ```python
      layer = LayerNorm(...)
      array = jax.vmap(layer)(array)
      ```
docstr: **Arguments:**
  
  - `shape`: Shape of the input.
  - `eps`: Value added to denominator for numerical stability.
  - `use_weight`: Whether the module has learnable affine weights.
  - `use_bias`: Whether the module has learnable affine biases.
  - `elementwise_affine`: Deprecated alternative to `use_weight` and `use_bias`.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","shape","`",":","Shape","of","the","input",".","-","`","eps","`",":","Value","added","to","denominator","for","numerical","stability",".","-","`","use_weight","`",":","Whether","the","module","has","learnable","affine","weights",".","-","`","use_bias","`",":","Whether","the","module","has","learnable","affine","biases",".","-","`","elementwise_affine","`",":","Deprecated","alternative","to","`","use_weight","`","and","`","use_bias","`","."]
code: def __init__(
      self,
      shape: Union[int, Sequence[int]],
      eps: float = 1e-5,
      use_weight: bool = True,
      use_bias: bool = True,
      *,
      elementwise_affine: Optional[bool] = None,
  ):
      """**Arguments:**
  
      - `shape`: Shape of the input.
      - `eps`: Value added to denominator for numerical stability.
      - `use_weight`: Whether the module has learnable affine weights.
      - `use_bias`: Whether the module has learnable affine biases.
      - `elementwise_affine`: Deprecated alternative to `use_weight` and `use_bias`.
      """
      if isinstance(shape, int):
          shape = (shape,)
      else:
          shape = tuple(shape)
      self.shape = shape
      self.eps = eps
      if elementwise_affine is not None:
          use_weight = elementwise_affine
          use_bias = elementwise_affine
          warnings.warn(
              "LayerNorm(elementwise_affine=...) is deprecated "
              "in favour of LayerNorm(use_weight=...) and LayerNorm(use_bias=...)"
          )
      self.use_weight = use_weight
      self.use_bias = use_bias
      self.weight = jnp.ones(shape) if use_weight else None
      self.bias = jnp.zeros(shape) if use_bias else None
  
imports: import functools
  import warnings
  from collections.abc import Sequence
  from typing import Optional, overload, Union
  import jax
  import jax.numpy
  from jaxtyping import Array, Float, PRNGKeyArray
  from .._custom_types import sentinel
  from .._misc import left_broadcast_to
  from .._module import field, Module
  from ._stateful import State
label_desc: Use image node_id 1 to create a new LayerNorm object from inherited base classes: Module with example: obj = LayerNorm(shape, eps, use_weight, use_bias)