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:**
  
  - `x`: A JAX array, with the same shape as the `shape` passed to `__init__`.
  - `state`: Ignored; provided for interchangability with the
      [`equinox.nn.BatchNorm`][] API.
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
      (Keyword only argument.)
  
  **Returns:**
  
  The output is a JAX array of the same shape as `x`.
  
  If `state` is passed, then a 2-tuple of `(output, state)` is returned. The state
  is passed through unchanged. If `state` is not passed, then just the output is
  returned.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","x","`",":","A","JAX","array",",","with","the","same","shape","as","the","`","shape","`","passed","to","`","__init__","`",".","-","`","state","`",":","Ignored",";","provided","for","interchangability","with","the","[","`","equinox.nn.BatchNorm","`","]","[","]","API",".","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","(","Keyword","only","argument",".",")","*","*","Returns",":","*","*","The","output","is","a","JAX","array","of","the","same","shape","as","`","x","`",".","If","`","state","`","is","passed",",","then","a","2-tuple","of","`","(","output",",","state",")","`","is","returned",".","The","state","is","passed","through","unchanged",".","If","`","state","`","is","not","passed",",","then","just","the","output","is","returned","."]
code: def __call__(
      self,
      x: Float[Array, "*shape"],
      state: State = sentinel,
      *,
      key: Optional[PRNGKeyArray] = None,
  ) -> Union[Array, tuple[Array, State]]:
      """**Arguments:**
  
      - `x`: A JAX array, with the same shape as the `shape` passed to `__init__`.
      - `state`: Ignored; provided for interchangability with the
          [`equinox.nn.BatchNorm`][] API.
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
          (Keyword only argument.)
  
      **Returns:**
  
      The output is a JAX array of the same shape as `x`.
  
      If `state` is passed, then a 2-tuple of `(output, state)` is returned. The state
      is passed through unchanged. If `state` is not passed, then just the output is
      returned.
      """
      if x.shape != self.shape:
          raise ValueError(
              "`RMSNorm(shape)(x)` must satisfy the invariant `shape == x.shape`"
              f"Received `shape={self.shape} and `x.shape={x.shape}`. You might need "
              "to replace `rms_norm(x)` with `jax.vmap(rms_norm)(x)`.\n"
          )
      inv_rms = jax.lax.rsqrt(jnp.mean(x**2) + self.eps)
      out = inv_rms * x
      if self.use_weight:
          out = self.weight * out
      if self.use_bias:
          out = out + self.bias
      if state is sentinel:
          return out
      else:
          return out, 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 4 for calling the RMSNorm obj's underlying member method code with example usage: obj.__call__(x, state) and returns: out, out,  state