Skip to content

tenzr::nn::RMSNorm

Root-mean-square normalization module.

#include <norm.h>

Inherits from tenzr::nn::Module

Public Functions

Name
RMSNorm(int64_t dim, double eps, DType dtype)
template <typename T =float>
RMSNorm(int64_t dim, double eps =1.0e-6)
Construct with the parameter dtype inferred from T.
template <typename T >
RMSNorm(std::type_identity< T > , int64_t dim, double eps =1.0e-6)
Construct from an explicit std::type_identity dtype tag.
virtual Tensor Forward(const Tensor & input) override
Tensor & Weight()
const Tensor & Weight() const
int64_t Dim() const
double Eps() const
virtual std::string Name() const override
template <typename T =float>
RMSNorm
Create(int64_t dim, double eps =1.0e-6)
Static factory mirroring the templated constructor.

Additional inherited members

Public Functions inherited from tenzr::nn::Module

Name
virtual ~Module() =default
Tensor operator()(const Tensor & input)
virtual void Train(bool mode)
void Train()
Switch to training mode.
void Eval()
Switch to evaluation mode.
bool IsTraining() const
Tensor & RegisterParameter(std::string name, Tensor param)
template <typename T ,typename... Args>
std::shared_ptr< T >
RegisterSubmodule(std::string name, Args &&... args)
std::shared_ptr< Module > RegisterSubmodule(std::string name, std::shared_ptr< Module > submod)
std::vector< Tensor > Parameters(bool recurse =true) const
std::unordered_map< std::string, Tensor > NamedParameters(std::string prefix ="", bool recurse =true) const
std::unordered_map< std::string, Tensor > StateDict() const
void LoadStateDict(const std::unordered_map< std::string, Tensor > & state)

Protected Attributes inherited from tenzr::nn::Module

Name
std::unordered_map< std::string, Tensor > parameters_
std::unordered_map< std::string, std::shared_ptr< Module > > submodules_
bool training_

Public Functions Documentation

function RMSNorm

RMSNorm(
    int64_t dim,
    double eps,
    DType dtype
)

Parameters:

  • dim Size of the normalized dimension.
  • eps Variance epsilon.
  • dtype Parameter dtype.
  • dim Size of the normalized dimension.
  • eps Stabilizing epsilon.
  • dtype Parameter dtype.

Construct with an explicit runtime dtype. Create a root-mean-square normalization over the last dimension.

function RMSNorm

template <typename T  =float>
inline explicit RMSNorm(
    int64_t dim,
    double eps =1.0e-6
)

Construct with the parameter dtype inferred from T.

function RMSNorm

template <typename T >
inline explicit RMSNorm(
    std::type_identity< T > ,
    int64_t dim,
    double eps =1.0e-6
)

Construct from an explicit std::type_identity dtype tag.

function Forward

virtual Tensor Forward(
    const Tensor & input
) override

Parameters:

  • input Input tensor.
  • input Input tensor whose last dimension equals dim.

Return:

  • Normalized tensor with the same shape.
  • Normalized tensor with the same shape.

Reimplements: tenzr::nn::Module::Forward

Normalize input by its root-mean-square. Scale the last dimension by the inverse root-mean-square.

function Weight

inline Tensor & Weight()

Return: Mutable reference to the scale weight.

function Weight

inline const Tensor & Weight() const

Return: Const reference to the scale weight.

function Dim

inline int64_t Dim() const

Return: Size of the normalized dimension.

function Eps

inline double Eps() const

Return: Variance epsilon.

function Name

inline virtual std::string Name() const override

Return: "RMSNorm".

Reimplements: tenzr::nn::Module::Name

function Create

template <typename T  =float>
static inline RMSNorm Create(
    int64_t dim,
    double eps =1.0e-6
)

Static factory mirroring the templated constructor.