ClikaRT::nn::RMSNorm

class

Header: ClikaRT/nn/rms_norm.h
Inherits: ClikaRT::nn::Module

A bound RMSNorm, the public face of the runtime's ops::rms_norm, exposed as an nn::Module leaf. Unlike the stateless ops::rms_norm free function, an RMSNorm binds its weight (gain) / bias ONCE and materializes them into contiguous layout at initialize(), so every forward reuses the bound params; build one per norm site at load, reuse it every step.

The gain / bias stay in their OWN dtype (a fp32 gain against a bf16 activation is read at full precision, never cast at rest); initialize() does layout work only. The output dtype follows the input x.

The lifecycle (uniform across every weight-bearing nn module): make(normalized_size, ...) registers storage-free weight / bias slots under their canonical names; set_weights(...) (positional) or load_state_dict(...) (a checkpoint, by dotted name) binds them; the FIRST forward materializes contiguous (once, thread-safe). initialize() remains an optional warm-up.

Held via std::shared_ptr (an nn::Module leaf); copy/move are pinned by the base.

Root-mean-square normalization module.

$$y_i = \frac{x_i}{\sqrt{\tfrac{1}{n}\sum_j x_j^2 + \varepsilon}} \cdot \gamma_i$$

The RMS is taken over the trailing normalized_shape dims, the module form of ops::rms_norm, holding gamma as its weight. The module owns its weights: construct with make(...), or declare shapes and bind a checkpoint via load_state_dict. After the first forward (or initialize()) the weight lives ONLY in the backend's packed form (one resident copy); to(dtype) restores, casts, and repacks on the next forward.

auto norm = ClikaRT::nn::RMSNorm::make(gamma);
auto y    = norm->forward(x);

Static member functions

make()

static std::shared_ptr<RMSNorm> make(
    std::int64_t normalized_size,
    bool bias = true,
    double eps = 1e-6,
    std::optional<ops::Activation> activation = std::nullopt,
    DataType dtype = DataType::Float32,
    Device device = Device::cpu()
)

Construct from config with storage-free slots: registers a shape-only weight [normalized_size] (and, with bias, a bias [normalized_size]) parameter on dtype/device; no storage is allocated. named_parameters() reports the slots, so a strict load_state_dict both expects and satisfies them; the first forward builds. Raises ClikaRT::Error on a bad normalized size.

Declared in ClikaRT/nn/rms_norm.h, line 63

Member functions

set_weights()

void set_weights(Tensor weight, OptionalTensor bias = {})

Bind the declared slots positionally: weight (gain, [normalized_size]) and, when the module was made with one, bias. The declared placement wins; shapes must match the declaration. Re-binding drops the built primitive; the next forward rebuilds. Raises ClikaRT::Error on a geometry mismatch or a bias without a declared bias slot.

Declared in ClikaRT/nn/rms_norm.h, line 75

~RMSNorm()

~RMSNorm() override

Declared in ClikaRT/nn/rms_norm.h, line 79

initialize_impl()

virtual Result<void> initialize_impl() override

Optional warm-up: materialize the bound weight / bias slots contiguous NOW (layout-only, dtype-preserving; idempotent, thread-safe). Every declared slot must hold a real (loaded) tensor; a still-fake slot is a clean error.

Declared in ClikaRT/nn/rms_norm.h, line 85

to_impl(StreamOrDevice)

virtual Result<void> to_impl(StreamOrDevice where) override

Move to a placement (Device / Stream) or cast to a dtype. Rebuilds the bound primitive on the target after moving the registry slots.

Declared in ClikaRT/nn/rms_norm.h, line 89

to_impl(DataType)

virtual Result<void> to_impl(DataType dtype) override

Declared in ClikaRT/nn/rms_norm.h, line 90

forward()

Tensor forward(Tensor x) const

out = ACT( x · rsqrt(mean(x²) + eps) · weight + bias ); x [*, N] -> [*, N], output dtype = x's. The first call builds (once, thread-safe); later calls reuse the bound primitive. Raises ClikaRT::Error on a shape/dtype mismatch or a still-fake slot.

Declared in ClikaRT/nn/rms_norm.h, line 97