class_docstr: A simplified version of LayerNorm which rescales the inputs, but does not center
  them. Optionally applies a learned reweighting of the transformed array afterward.
  
  Given an input array $x$, this layer computes
  
  $$\frac{x}{\sqrt{\varepsilon + \frac{1}{n}\Vert x \Vert^2_2}} \gamma + \beta$$
  
  where $\Vert x \Vert^2_2 = \sum_{i=1}^n x_i^2$, $n = \dim(x)$, and $\gamma$ is a
  learned array with the same shape as $x$ if `use_weight=True`, or
  $\gamma = 1$ if `use_weight=False`, as proposed in
  [this paper](https://browse.arxiv.org/abs/2307.14995). `\beta` is an optional bias
  term.
  
  ??? cite
  
      [Root Mean Square Layer Normalization](https://browse.arxiv.org/abs/1910.07467)
  
      ```bibtex
      @article{zhang2019root,
          title={Root Mean Square Layer Normalization},
          author={Biao Zhang and Rico Sennrich},
          year={2019},
          journal={arXiv:1910.07467}
      }
      ```
code: def __call__(
      self,
      x: Array,
      state: State,
      *,
      key: Optional[PRNGKeyArray] = None,
  ) -> tuple[Array, State]:
      ...
  
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 3 for calling the RMSNorm obj's underlying member method code with example usage: obj.__call__(x, state) without return types