vyre-libs 0.6.1

vyre Category A library ecosystem - pure-IR compositions over vyre-ops hardware primitives
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

//! Shared RMS expression builders for normalization and fused linear kernels.

use vyre::ir::{Expr, UnOp};

pub(crate) const EMPTY_RMS_FIX: &str =
    "Fix: rms_norm n=0 is invalid; pass at least one element or bypass normalization.";

pub(crate) fn square_expr(value: Expr) -> Expr {
    Expr::mul(value.clone(), value)
}

pub(crate) fn inverse_rms_expr(sum_sq: Expr, n: u32, eps: f32) -> Expr {
    Expr::UnOp {
        op: UnOp::InverseSqrt,
        operand: Box::new(Expr::add(
            Expr::div(sum_sq, Expr::f32(n as f32)),
            Expr::f32(eps),
        )),
    }
}