Struct RmsNorm 

pub struct RmsNorm<D: Dtype, IT, OT, const RANK: usize> {
    pub normalized_shape: Vec<usize>,
    pub eps: f64,
    pub affine: bool,
    /* private fields */
}

Root Mean Square normalisation.

Normalises over the last normalized_shape.len() dimensions by dividing by the RMS (no mean subtraction): y = x / rms(x) * γ

Output shape equals input shape.

Parameters:

• normalized_shape — axes to normalise over
• eps — numerical stability constant (default 1e-8)
• affine — if true, learns per-element γ (default true)

Fields

normalized_shape: Vec<usize>

The shape of the trailing axes to normalize over.

eps: f64

Numerical stability constant.

affine: bool

Whether to learn a per-element scale (gamma) parameter.

Implementations

impl<D: Dtype, IT, OT, const RANK: usize> RmsNorm<D, IT, OT, RANK>

pub fn new(normalized_shape: impl Into<Vec<usize>>) -> Self

Creates a new layer with default eps, affine on.

pub fn with_eps(self, eps: f64) -> Self

Sets the numerical stability constant.

pub fn with_affine(self, affine: bool) -> Self

Sets whether to learn a per-element scale parameter.

Trait Implementations

impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT> for RmsNorm<D, IT, OT, RANK>

type Output = OT

This layer’s output type.

fn call(&self, _input: IT) -> Self::Output

Applies this layer to input.

impl<D: Dtype, const RANK: usize> Layer<SymTensor> for RmsNorm<D, SymTensor, SymTensor, RANK>

type Output = SymTensor

This layer’s output type.

fn call(&self, input: SymTensor) -> SymTensor

Applies this layer to input.