Skip to main content

add_rms_norm

Function add_rms_norm 

Source
pub fn add_rms_norm<'a>(
    x: &Array,
    residual: &Array,
    weight: impl Into<Option<&'a Array>>,
    eps: f32,
) -> (Array, Array)
Expand description

A residual connection followed by RMS normalisation: h = x + residual and rms_norm(h, weight, eps), returned as (h, normed), from one GPU kernel. The results equal x + residual then rms_norm exactly; a weight that would widen the sum’s dtype runs the two operations.