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: Constructs a RMSNorm module.
  
  Args:
    axis: Integer, list of integers, or slice indicating which axes to
      normalize over.
    eps: Small epsilon to avoid division by zero variance. Defaults to 1e-5.
    scale_init: Optional initializer for gain (aka scale). By default, one.
    name: The module name.
    create_scale: Bool, defines whether to create a trainable scale
      per channel applied after the normalization.
    param_axis: Axis used to determine the parameter shape of the learnable
      scale/offset. Sonnet sets this to the channel/feature axis (e.g. to
      ``-1`` for ``NHWC``). Other libraries set this to the same as the
      reduction axis (e.g. ``axis=param_axis``). `None` defaults to (-1,).
docstr_tok: ["Constructs","a","RMSNorm","module",".","Args",":","axis",":","Integer",",","list","of","integers",",","or","slice","indicating","which","axes","to","normalize","over",".","eps",":","Small","epsilon","to","avoid","division","by","zero","variance",".","Defaults","to","1e-5",".","scale_init",":","Optional","initializer","for","gain","(","aka","scale",")",".","By","default",",","one",".","name",":","The","module","name",".","create_scale",":","Bool",",","defines","whether","to","create","a","trainable","scale","per","channel","applied","after","the","normalization",".","param_axis",":","Axis","used","to","determine","the","parameter","shape","of","the","learnable","scale\/offset",".","Sonnet","sets","this","to","the","channel\/feature","axis","(","e.g",".","to","``","-1","``","for","``","NHWC","``",")",".","Other","libraries","set","this","to","the","same","as","the","reduction","axis","(","e.g",".","``","axis=param_axis","``",")",".","`","None","`","defaults","to","(","-1",",",")","."]
code: def __init__(
      self,
      axis: AxisOrAxes,
      eps: float = 1e-5,
      scale_init: Optional[hk.initializers.Initializer] = None,
      name: Optional[str] = None,
      create_scale: bool = True,
      *,
      param_axis: Optional[AxisOrAxes] = None,
  ):
      """Constructs a RMSNorm module.
  
      Args:
        axis: Integer, list of integers, or slice indicating which axes to
          normalize over.
        eps: Small epsilon to avoid division by zero variance. Defaults to 1e-5.
        scale_init: Optional initializer for gain (aka scale). By default, one.
        name: The module name.
        create_scale: Bool, defines whether to create a trainable scale
          per channel applied after the normalization.
        param_axis: Axis used to determine the parameter shape of the learnable
          scale/offset. Sonnet sets this to the channel/feature axis (e.g. to
          ``-1`` for ``NHWC``). Other libraries set this to the same as the
          reduction axis (e.g. ``axis=param_axis``). `None` defaults to (-1,).
      """
      super().__init__(name=name)
      if not create_scale and scale_init is not None:
          raise ValueError(
              "Cannot set `scale_init` if `create_scale=False`."
          )
      if isinstance(axis, slice):
          self.axis = axis
      elif isinstance(axis, int):
          self.axis = (axis,)
      elif isinstance(axis, abc.Iterable) and all(
          isinstance(ax, int) for ax in axis
      ):
          self.axis = tuple(axis)
      else:
          raise ValueError(
              "`axis` should be an int, slice or iterable of ints."
          )
  
      self.eps = eps
      self.create_scale = create_scale
      self.scale_init = scale_init or jnp.ones
      if param_axis is None:
          self.param_axis = (-1,)
      else:
          self.param_axis = layer_norm.to_axes_or_slice(param_axis)
  
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 1 to create a new RMSNorm object from inherited base classes: hk with example: obj = RMSNorm(axis, eps, scale_init, name, create_scale)