zsfm_burn/
norm.rs

1use burn::tensor::{activation, Tensor};
2#[cfg(test)]
3use burn::tensor::{Device, TensorData};
4
5/// RMSNorm over the last dim, with weight.
6pub fn burn_rms_norm(x: Tensor<2>, w: Tensor<1>, eps: f32) -> Tensor<2> {
7    let mean_sq = x.clone().square().mean_dim(1);
8    let rms = (mean_sq.add_scalar(eps)).sqrt();
9    let normed = x / rms;
10    normed * w.unsqueeze_dim::<2>(0)
11}
12
13/// LayerNorm over the last dim, with weight and bias.
14pub fn burn_layer_norm(x: Tensor<2>, w: Tensor<1>, b: Tensor<1>, eps: f32) -> Tensor<2> {
15    let mean = x.clone().mean_dim(1);
16    let centered = x - mean;
17    let var = centered.clone().square().mean_dim(1);
18    let std = (var.add_scalar(eps)).sqrt();
19    let normed = centered / std;
20    normed * w.unsqueeze_dim::<2>(0) + b.unsqueeze_dim::<2>(0)
21}
22
23/// Softmax over the last dim for host-checked logits.
24pub fn burn_softmax_last_dim(x: Tensor<2>) -> Tensor<2> {
25    activation::softmax(x, 1)
26}
27
28/// RMSNorm over the last dim for rank-D input.
29pub fn burn_rms_norm_nd<const D: usize>(x: Tensor<D>, w: Tensor<1>, eps: f32) -> Tensor<D> {
30    let dims = x.dims();
31    let last = dims[D - 1];
32    let lead: usize = dims[..D - 1].iter().product();
33    burn_rms_norm(x.reshape([lead, last]), w, eps).reshape(dims)
34}
35
36/// LayerNorm over the last dim for rank-D input.
37pub fn burn_layer_norm_nd<const D: usize>(
38    x: Tensor<D>,
39    w: Tensor<1>,
40    b: Tensor<1>,
41    eps: f32,
42) -> Tensor<D> {
43    let dims = x.dims();
44    let last = dims[D - 1];
45    let lead: usize = dims[..D - 1].iter().product();
46    burn_layer_norm(x.reshape([lead, last]), w, b, eps).reshape(dims)
47}
48
49/// LayerNorm without affine parameters over the last dim (2D).
50pub fn burn_layer_norm_no_affine(x: Tensor<2>, eps: f32) -> Tensor<2> {
51    let mean = x.clone().mean_dim(1);
52    let centered = x - mean;
53    let var = centered.clone().square().mean_dim(1);
54    let std = (var.add_scalar(eps)).sqrt();
55    centered / std
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61
62    fn dev() -> Device {
63        Device::flex()
64    }
65
66    #[test]
67    fn rms_matches_manual() {
68        let device = dev();
69        let x = Tensor::<2>::from_data(TensorData::from([[1.0f32, 2.0, 3.0, -4.0]]), &device);
70        let w = Tensor::<1>::from_data(TensorData::from([2.0f32, 2.0, 2.0, 2.0]), &device);
71        let got: Vec<f32> = burn_rms_norm(x, w, 1e-6).to_data().try_to_vec().unwrap();
72        let vals = [1.0f32, 2.0, 3.0, -4.0];
73        let mean_sq: f32 = vals.iter().map(|v| v * v).sum::<f32>() / 4.0;
74        let rms = (mean_sq + 1e-6).sqrt();
75        for (g, v) in got.iter().zip(vals.iter()) {
76            assert!((g - v / rms * 2.0).abs() < 1e-5, "{g} vs {v}");
77        }
78    }
79
80    #[test]
81    fn softmax_sums_to_one() {
82        let device = dev();
83        let x = Tensor::<2>::from_data(TensorData::from([[1.0f32, 2.0, 3.0]]), &device);
84        let p: Vec<f32> = burn_softmax_last_dim(x).to_data().try_to_vec().unwrap();
85        assert!((p.iter().sum::<f32>() - 1.0).abs() < 1e-5);
86    }
87}
zsfm_burn/norm.rs

zsfm_burn/
norm.rs
