RMSNorm

RMSNorm#

class seli.net.RMSNorm(eps: float = 1e-06, offset: float | int = 1)[source]#

Bases: Module

Scale the input by the reciprocal of the root mean square along the last axis. Then add a learnable offset and scale by a learnable weight along the last axis.

Parameters:

  • eps (float) – Epsilon value for numerical stability.

  • axis (int) – The axis to calculate the root mean square.

Methods Summary

__call__(x)

Call self as a function.

Methods Documentation

__call__(x: Float[Array, '*batch dim']) → Float[Array, '*batch dim'][source]#

Call self as a function.

Contents