class_docstr: RMSNorm module.
  
  RMSNorm provides an alternative that can be both faster and more stable than
  LayerNorm. The inputs are normalized by the root-mean-squared (RMS) and scaled
  by a learned parameter, but they are not recentered around their mean.
  
  See https://arxiv.org/pdf/1910.07467.pdf
docstr: Connects the layer norm.
  
  Args:
    inputs: An array, where the data format is ``[N, ..., C]``.
  
  Returns:
    The normalized array, of the same shape as the inputs.
docstr_tok: ["Connects","the","layer","norm",".","Args",":","inputs",":","An","array",",","where","the","data","format","is","``","[","N",",","...",",","C","]","``",".","Returns",":","The","normalized","array",",","of","the","same","shape","as","the","inputs","."]
code: def __call__(self, inputs: jax.Array):
      """Connects the layer norm.
  
      Args:
        inputs: An array, where the data format is ``[N, ..., C]``.
  
      Returns:
        The normalized array, of the same shape as the inputs.
      """
      axis = self.axis
      if isinstance(axis, slice):
          axis = tuple(range(inputs.ndim)[axis])
  
      param_axis = layer_norm.to_abs_axes(self.param_axis, inputs.ndim)
      if param_axis == (inputs.ndim - 1,):
          # For param_axis=-1 we store non-broadcast param shape for compatibility
          # with older checkpoints.
          param_shape = inputs.shape[-1:]
      else:
          param_shape = tuple(
              (inputs.shape[i] if i in param_axis else 1)
              for i in range(inputs.ndim)
          )
      if self.create_scale:
          scale = hk.get_parameter(
              "scale", param_shape, inputs.dtype, init=self.scale_init
          )
          scale = jnp.broadcast_to(scale, inputs.shape)
      else:
          scale = 1.0
  
      mean_squared = jnp.mean(
          jnp.square(inputs), axis=axis, keepdims=True
      )
      mean_squared = jnp.broadcast_to(mean_squared, inputs.shape)
  
      return inputs * scale * jax.lax.rsqrt(mean_squared + self.eps)
  
file_path: dm-haiku/haiku/_src/rms_norm.py
imports: from collections import abc
  from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import layer_norm
  from haiku._src import module
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the RMSNorm obj's underlying member method code with example usage: obj.__call__(inputs) and returns: unknown