Skip to main content

tract_linalg/generic/
sigmoid.rs

1#![allow(clippy::excessive_precision)]
2use crate::frame::element_wise::ElementWiseKer;
3use tract_data::internal::*;
4
5/// f32 sigmoid, as a rational minimax fit of `sigmoid(x) - 0.5` over `[LOW, HIGH]`.
6///
7/// The sum is clamped to `[0, 1]` to keep the range consumers rely on: on either tail
8/// `p / q` approaches ±0.5 and the final `+ 0.5` cancels down to ~1e-8, below the ~6e-8
9/// f32 rounding error of `p / q` itself, so the sum can land just outside the interval.
10/// NaN still propagates.
11pub fn ssigmoid(x: f32) -> f32 {
12    const LOW: f32 = -18.6;
13    const HIGH: f32 = -LOW;
14
15    const ALPHA_13: f32 = -4.433153405e-18;
16    const ALPHA_11: f32 = 1.169974371e-14;
17    const ALPHA_9: f32 = -1.875289645e-11;
18    const ALPHA_7: f32 = 4.257889523e-8;
19    const ALPHA_5: f32 = 0.00004811817576;
20    const ALPHA_3: f32 = 0.008163842030;
21    const ALPHA_1: f32 = 0.2499999971;
22    const BETA_6: f32 = 3.922935744e-6;
23    const BETA_4: f32 = 0.001524872358;
24    const BETA_2: f32 = 0.1159886749;
25    const BETA_0: f32 = 1.0;
26
27    let x = x.clamp(LOW, HIGH);
28
29    let x2 = x * x;
30
31    let p = ALPHA_13;
32    let p = x2 * p + ALPHA_11;
33    let p = x2 * p + ALPHA_9;
34    let p = x2 * p + ALPHA_7;
35    let p = x2 * p + ALPHA_5;
36    let p = x2 * p + ALPHA_3;
37    let p = x2 * p + ALPHA_1;
38    let p = p * x;
39
40    let q = BETA_6;
41    let q = x2 * q + BETA_4;
42    let q = x2 * q + BETA_2;
43    let q = x2 * q + BETA_0;
44
45    (p / q + 0.5).clamp(0.0, 1.0)
46}
47
48/// f16 sigmoid, as a rational minimax fit of `sigmoid(x) - 0.5` over `[LOW, HIGH]`.
49///
50/// Needs no output clamp, unlike [`ssigmoid`]: the `+ 0.5` cancellation stays inside
51/// (0, 1) for every non-NaN f16 input.
52pub fn hsigmoid(x: f16) -> f16 {
53    /*
54     * (x (0.249895 + x^2 (0.00400222 - 0.0000124702 x^2)))
55     * /
56     * (1. + 0.098734 x^2)
57     */
58
59    const LOW: f16 = f16::from_f32_const(-6.92);
60    const HIGH: f16 = f16::from_f32_const(6.92);
61
62    const ALPHA_5: f16 = f16::from_f32_const(-0.0000124702);
63    const ALPHA_3: f16 = f16::from_f32_const(0.00400222);
64    const ALPHA_1: f16 = f16::from_f32_const(0.249895);
65
66    const BETA_2: f16 = f16::from_f32_const(0.098734);
67    const BETA_0: f16 = f16::from_f32_const(1.0);
68
69    let x = x.clamp(LOW, HIGH);
70
71    let x2 = x * x;
72
73    let p = ALPHA_5;
74    let p = x2 * p + ALPHA_3;
75    let p = x2 * p + ALPHA_1;
76    let p = p * x;
77
78    let q = BETA_2;
79    let q = x2 * q + BETA_0;
80
81    p / q + f16::from_f32_const(0.5)
82}
83
84#[derive(Clone, Debug)]
85pub struct SSigmoid4;
86
87impl ElementWiseKer<f32> for SSigmoid4 {
88    fn name() -> &'static str {
89        "generic"
90    }
91
92    fn alignment_bytes() -> usize {
93        16
94    }
95
96    fn alignment_items() -> usize {
97        4
98    }
99
100    fn nr() -> usize {
101        4
102    }
103
104    fn run(x: &mut [f32], _: ()) {
105        debug_assert!(x.len() % Self::nr() == 0);
106        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
107        x.iter_mut().for_each(|px| *px = ssigmoid(*px))
108    }
109}
110
111#[derive(Clone, Debug)]
112pub struct HSigmoid8;
113
114impl ElementWiseKer<f16> for HSigmoid8 {
115    fn name() -> &'static str {
116        "generic"
117    }
118
119    fn alignment_bytes() -> usize {
120        16
121    }
122
123    fn alignment_items() -> usize {
124        4
125    }
126
127    fn nr() -> usize {
128        8
129    }
130
131    fn run(x: &mut [f16], _: ()) {
132        debug_assert!(x.len() % Self::nr() == 0);
133        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
134        x.iter_mut().for_each(|px| *px = hsigmoid(*px))
135    }
136}
137
138#[cfg(test)]
139#[macro_use]
140pub mod s {
141    sigmoid_frame_tests!(true, f32, crate::generic::sigmoid::SSigmoid4);
142}
143
144#[cfg(test)]
145#[macro_use]
146pub mod h {
147    sigmoid_frame_tests!(true, tract_data::internal::f16, crate::generic::sigmoid::HSigmoid8);
148}