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}
      }
      ```
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 shift.
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","shift","."]
code: def __init__(
      self,
      shape: Union[int, Sequence[int]],
      eps: float = 1e-5,
      use_weight: bool = True,
      use_bias: bool = True,
  ):
      """**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 shift.
      """
      if isinstance(shape, int):
          shape = (shape,)
      else:
          shape = tuple(shape)
      self.shape = shape
      self.eps = eps
      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 RMSNorm object from inherited base classes: Module with example: obj = RMSNorm(shape, eps, use_weight, use_bias)