combs_models/
norm.rs

1//! RMSNorm (Llama-style, no learnable bias).
2
3use burn::tensor::{Tensor, backend::Backend};
4
5use crate::precision::{to_f32, to_float};
6
7/// `y = x / rms(x) * w` where `rms` is taken over the last dimension and
8/// `eps` is added inside the square root.
9///
10/// The reduction runs in f32 for f16 stability (no-op in f32 builds).
11pub fn rms_norm<B: Backend, const D: usize>(
12    x: Tensor<B, D>,
13    weight: Tensor<B, 1>,
14    eps: f64,
15) -> Tensor<B, D> {
16    let out_dtype = x.dtype();
17    let dims = x.dims();
18    let hidden = dims[D - 1];
19
20    let xf = to_f32(x);
21    // mean(x^2) over the last dim, keeping rank for broadcasting.
22    let mean_sq = xf.clone().powf_scalar(2.0).mean_dim(D - 1);
23    let inv_rms = mean_sq.add_scalar(eps).sqrt().recip();
24
25    let mut shape = [1usize; D];
26    shape[D - 1] = hidden;
27    let y = xf * inv_rms * to_f32(weight).reshape(shape);
28    to_float(y, out_dtype)
29}
30
31/// Gemma-style RMSNorm: `y = x / rms(x) * (1 + w)` — the learnable weight
32/// is zero-centered (HF `Gemma3RMSNorm`).
33pub fn gemma_rms_norm<B: Backend, const D: usize>(
34    x: Tensor<B, D>,
35    weight: Tensor<B, 1>,
36    eps: f64,
37) -> Tensor<B, D> {
38    rms_norm(x, weight.add_scalar(1.0), eps)
39}
40
41/// LayerNorm (SigLIP-style, learnable weight + bias):
42/// `y = (x - μ) / sqrt(σ² + eps) * w + b`, statistics over the last dim.
43pub fn layer_norm<B: Backend, const D: usize>(
44    x: Tensor<B, D>,
45    weight: Tensor<B, 1>,
46    bias: Tensor<B, 1>,
47    eps: f64,
48) -> Tensor<B, D> {
49    let dims = x.dims();
50    let hidden = dims[D - 1];
51
52    let out_dtype = x.dtype();
53    let xf = to_f32(x);
54    let mean = xf.clone().mean_dim(D - 1);
55    let centered = xf - mean;
56    let var = centered.clone().powf_scalar(2.0).mean_dim(D - 1);
57    let inv_std = var.add_scalar(eps).sqrt().recip();
58
59    let mut shape = [1usize; D];
60    shape[D - 1] = hidden;
61    let y = centered * inv_std * to_f32(weight).reshape(shape.clone())
62        + to_f32(bias).reshape(shape);
63    to_float(y, out_dtype)
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use burn::tensor::TensorData;
70
71    type TestBackend = burn::backend::NdArray<f32>;
72
73    #[test]
74    fn normalizes_rows_to_unit_rms() {
75        let device = burn::tensor::Device::<TestBackend>::default();
76        let x: Tensor<TestBackend, 2> = Tensor::from_data(
77            TensorData::new(vec![3.0f32, 4.0, 1.0, -2.0, 0.5, 2.5], [2, 3]),
78            &device,
79        );
80        let w: Tensor<TestBackend, 1> = Tensor::ones([3], &device);
81        let y = rms_norm(x, w, 1e-6);
82        // With unit weights, each row of y must have RMS == 1 (up to eps).
83        let rms = y
84            .powf_scalar(2.0)
85            .mean_dim(1)
86            .sqrt()
87            .into_data()
88            .to_vec::<f32>()
89            .unwrap();
90        for (i, r) in rms.iter().enumerate() {
91            assert!((r - 1.0).abs() < 1e-4, "row {i} rms = {r}");
92        }
93    }
94
95    #[test]
96    fn applies_weight() {
97        let device = burn::tensor::Device::<TestBackend>::default();
98        let x: Tensor<TestBackend, 2> =
99            Tensor::from_data(TensorData::new(vec![1.0f32, 2.0], [1, 2]), &device);
100        let w: Tensor<TestBackend, 1> =
101            Tensor::from_data(TensorData::new(vec![2.0f32, 2.0], [2]), &device);
102        let y = rms_norm(x.clone(), w, 1e-6);
103        let z = rms_norm(x, Tensor::ones([2], &device), 1e-6);
104        let yv: Vec<f32> = y.into_data().to_vec().unwrap();
105        let zv: Vec<f32> = z.into_data().to_vec().unwrap();
106        for i in 0..2 {
107            assert!((yv[i] - 2.0 * zv[i]).abs() < 1e-4);
108        }
109    }
110}
combs_models/norm.rs

combs_models/
norm.rs
