axiom_engine 0.5.0

Axiom-TTT Inference Engine – Rust implementation using candle
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

use candle_core::{Result, Tensor, D};
use candle_nn::VarBuilder;

/// Root Mean Square Layer Normalization.
pub struct RMSNorm {
    weight: Tensor,
    eps: f32,
}

impl RMSNorm {
    pub fn new(dim: usize, eps: f32, vs: VarBuilder) -> Result<Self> {
        let weight = vs.get(dim, "weight")?;
        Ok(Self { weight, eps })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let variance = x.sqr()?.mean_keepdim(D::Minus1)?;
        let eps = Tensor::new(self.eps, x.device())?;
        let norm = x.broadcast_div(&variance.broadcast_add(&eps)?.sqrt()?)?;
        norm.broadcast_mul(&self.weight)
    }
}