kopitiam_tensor/tensor/
activation.rs

1//! `silu` and `gelu` activations.
2
3use kopitiam_core::{DType, Result};
4
5use crate::storage::Storage;
6
7use super::Tensor;
8
9/// `sqrt(2 / pi)`, the constant in the tanh approximation of GELU.
10const SQRT_2_OVER_PI: f32 = 0.797_884_6;
11
12impl Tensor {
13    /// SiLU / swish (Elfwing, Uchibe & Doya, 2017; also Hendrycks & Gimpel):
14    /// `x * sigmoid(x) = x / (1 + exp(-x))`. Used by LLaMA-family FFNs.
15    pub fn silu(&self) -> Result<Tensor> {
16        self.unary_f32(|x| x / (1.0 + (-x).exp()))
17    }
18
19    /// GELU (Hendrycks & Gimpel, 2016), tanh approximation:
20    /// `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`.
21    ///
22    /// This crate implements the tanh approximation rather than the exact
23    /// `0.5 * x * (1 + erf(x / sqrt(2)))` form because `erf` is not in
24    /// Rust's `std` — computing it exactly would mean shipping a rational
25    /// or Chebyshev `erf` approximation of our own, which is *more* code
26    /// and *more* numerical risk than the tanh form for a difference that
27    /// is below `1e-3` everywhere and is what GPT-2/GPT-NeoX-family models
28    /// were themselves trained and evaluated with (the tanh form is
29    /// sometimes called "gelu_new" in that lineage). If a model that
30    /// specifically requires the exact erf form shows up, add it as a
31    /// second method rather than replacing this one.
32    pub fn gelu(&self) -> Result<Tensor> {
33        self.unary_f32(|x| 0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044_715 * x * x * x)).tanh()))
34    }
35
36    fn unary_f32(&self, f: impl Fn(f32) -> f32) -> Result<Tensor> {
37        self.require_dtype(DType::F32)?;
38        let Storage::F32(data) = self.storage.as_ref() else { unreachable!() };
39        let out: Vec<f32> = self.logical_offsets().map(|i| f(data[i])).collect();
40        Tensor::from_f32(out, self.shape.clone())
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use kopitiam_core::Error;
47
48    use super::*;
49
50    fn assert_close(a: f32, b: f32, epsilon: f32) {
51        assert!((a - b).abs() < epsilon, "expected {b}, got {a}");
52    }
53
54    #[test]
55    fn silu_matches_hand_computation() {
56        // silu(0) = 0 * sigmoid(0) = 0 * 0.5 = 0
57        let t = Tensor::from_f32(vec![0.0, 1.0, -1.0], [3]).unwrap();
58        let out = t.silu().unwrap().to_vec_f32().unwrap();
59        assert_close(out[0], 0.0, 1e-6);
60        // silu(1) = 1 / (1 + e^-1) ~= 0.7310586
61        assert_close(out[1], 0.731_058_6, 1e-5);
62        // silu(-1) = -1 / (1 + e^1) ~= -0.2689414
63        assert_close(out[2], -0.268_941_4, 1e-5);
64    }
65
66    #[test]
67    fn silu_matches_the_reference_formula_across_a_range_of_inputs() {
68        // silu is *not* globally monotonic (it dips to a minimum around
69        // x ~= -1.278 before rising) so the meaningful correctness check is
70        // "matches x * sigmoid(x) pointwise", not a monotonicity claim.
71        let xs = vec![-5.0, -3.0, -1.278, -1.0, -0.5, 0.0, 0.5, 1.0, 3.0, 5.0];
72        let t = Tensor::from_f32(xs.clone(), [xs.len()]).unwrap();
73        let out = t.silu().unwrap().to_vec_f32().unwrap();
74        for (x, o) in xs.iter().zip(&out) {
75            let expected = x / (1.0 + (-x).exp());
76            assert_close(*o, expected, 1e-5);
77        }
78    }
79
80    #[test]
81    fn silu_is_monotonically_increasing_on_its_increasing_branch() {
82        // For x >= -1.278ish, silu is monotonically increasing; this is the
83        // range every real activation input in a trained model's residual
84        // stream overwhelmingly falls into after the first few layers.
85        let t = Tensor::from_f32(vec![-1.0, 0.0, 1.0, 2.0, 3.0], [5]).unwrap();
86        let out = t.silu().unwrap().to_vec_f32().unwrap();
87        for pair in out.windows(2) {
88            assert!(pair[1] > pair[0], "silu should be increasing here: {out:?}");
89        }
90    }
91
92    #[test]
93    fn gelu_matches_hand_computation_at_zero_and_matches_known_values() {
94        // gelu(0) = 0.
95        let t = Tensor::from_f32(vec![0.0, 1.0, -1.0], [3]).unwrap();
96        let out = t.gelu().unwrap().to_vec_f32().unwrap();
97        assert_close(out[0], 0.0, 1e-6);
98        // Known reference values for the tanh-approximation GELU (matches
99        // the widely used "gelu_new" implementation to ~1e-6).
100        assert_close(out[1], 0.841_192, 1e-4);
101        assert_close(out[2], -0.158_808, 1e-4);
102    }
103
104    #[test]
105    fn gelu_approaches_the_identity_for_large_positive_x_and_zero_for_large_negative_x() {
106        let t = Tensor::from_f32(vec![10.0, -10.0], [2]).unwrap();
107        let out = t.gelu().unwrap().to_vec_f32().unwrap();
108        assert_close(out[0], 10.0, 1e-3);
109        assert_close(out[1], 0.0, 1e-3);
110    }
111
112    #[test]
113    fn activations_reject_non_f32_input() {
114        let t = Tensor::from_i32(vec![1, 2, 3], [3]).unwrap();
115        assert!(matches!(t.silu(), Err(Error::DTypeMismatch { .. })));
116        assert!(matches!(t.gelu(), Err(Error::DTypeMismatch { .. })));
117    }
118}
kopitiam_tensor/tensor/activation.rs

kopitiam_tensor/tensor/
activation.rs
