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)
      ```
code: def __call__(
      self, x: Array, *, key: Optional[PRNGKeyArray] = None
  ) -> Array:
      ...
  
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 2 for calling the LayerNorm obj's underlying member method code with example usage: obj.__call__(x) without return types