tract_linalg/generic/
gelu.rs

1#![allow(clippy::excessive_precision)]
2use crate::generic::tanh::stanh;
3use tract_data::internal::*;
4
5// Tanh-form GELU approximation matching tract's GeluApproximate (pow=3, the
6// canonical Hendrycks-Gimpel/Open-AI form):
7//
8//     gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
9//
10// The fast variant (pow=2) is not exposed here; the graph op falls back to
11// scalar when fast_impl=true.
12
13const SQRT_2_OVER_PI: f32 = 0.7978845608028654;
14const COEF: f32 = 0.044715;
15
16routine_ew_rust!(generic;
17    f32,
18    generic_gelu_f32_4n,
19    4,
20    4,
21    fn run(x: &mut [f32], _: ()) {
22        debug_assert!(x.len() % Self::nr() == 0);
23        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
24        x.iter_mut().for_each(|px| {
25            let v = *px;
26            let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v);
27            *px = 0.5 * v * (1.0 + stanh(inner));
28        });
29    },
30    func(Gelu)
31);
32
33routine_ew_rust!(generic;
34    f16,
35    generic_gelu_f16_8n,
36    8,
37    8,
38    fn run(x: &mut [f16], _: ()) {
39        debug_assert!(x.len() % Self::nr() == 0);
40        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
41        x.iter_mut().for_each(|px| {
42            let v = px.to_f32();
43            let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v);
44            *px = f16::from_f32(0.5 * v * (1.0 + stanh(inner)));
45        });
46    },
47    func(Gelu)
48);
tract_linalg/generic/gelu.rs

tract_linalg/generic/
gelu.rs
