Skip to main content

tract_linalg/generic/
tanh.rs

1#![allow(clippy::excessive_precision)]
2use tract_data::internal::*;
3
4/// f32 tanh, as a rational minimax fit of `tanh` over `[LOW, HIGH]`.
5///
6/// The clamp keeps the `[-1, 1]` range consumers rely on without an output clamp: the two
7/// Horner chains and the division leave a few ulps of relative error, and the largest
8/// quotient any f32 input produces stays seven f32 steps under 1. It costs the tails their
9/// last few ulps — a saturated input returns `±(1 - 6.1e-7)`, not `±1`. NaN propagates.
10///
11/// Raising the clamp voids this: `1 - tanh(8.18)` is already under the rounding error, so
12/// the quotient crosses 1 at scattered inputs from there up.
13pub fn stanh(x: f32) -> f32 {
14    const LOW: f32 = -7.5;
15    const HIGH: f32 = 7.5;
16
17    const ALPHA_13: f32 = -8.488492677e-14;
18    const ALPHA_11: f32 = 5.277853000e-11;
19    const ALPHA_9: f32 = -2.022500419e-8;
20    const ALPHA_7: f32 = 0.00001115424833;
21    const ALPHA_5: f32 = 0.003103950131;
22    const ALPHA_3: f32 = 0.1308400453;
23    const ALPHA_1: f32 = 0.9999999934;
24
25    const BETA_6: f32 = 0.0002546136580;
26    const BETA_4: f32 = 0.02449515379;
27    const BETA_2: f32 = 0.4641733162;
28    const BETA_0: f32 = 1.0;
29
30    let x = x.clamp(LOW, HIGH);
31
32    let x2 = x * x;
33
34    let p = ALPHA_13;
35    let p = x2 * p + ALPHA_11;
36    let p = x2 * p + ALPHA_9;
37    let p = x2 * p + ALPHA_7;
38    let p = x2 * p + ALPHA_5;
39    let p = x2 * p + ALPHA_3;
40    let p = x2 * p + ALPHA_1;
41    let p = p * x;
42
43    let q = BETA_6;
44    let q = x2 * q + BETA_4;
45    let q = x2 * q + BETA_2;
46    let q = x2 * q + BETA_0;
47
48    p / q
49}
50
51/// f16 tanh, as a rational minimax fit of `tanh` over `[LOW, HIGH]`.
52///
53/// Needs no output clamp, like [`stanh`]: `1 - tanh(3.84)` is `9.3e-4`, about two
54/// f16 steps below 1, so the quotient keeps its margin for every f16 input.
55pub fn htanh(x: f16) -> f16 {
56    const LOW: f16 = f16::from_f32_const(-3.84);
57    const HIGH: f16 = f16::from_f32_const(3.84);
58
59    const ALPHA_3: f16 = f16::from_f32_const(0.082654955);
60    const ALPHA_1: f16 = f16::from_f32_const(0.99963124);
61
62    const BETA_4: f16 = f16::from_f32_const(0.0065383179);
63    const BETA_2: f16 = f16::from_f32_const(0.41401828);
64    const BETA_0: f16 = f16::from_f32_const(1.0);
65
66    let x = x.clamp(LOW, HIGH);
67
68    let x2 = x * x;
69
70    let p = ALPHA_3;
71    let p = x2 * p + ALPHA_1;
72    let p = p * x;
73
74    let q = BETA_4;
75    let q = x2 * q + BETA_2;
76    let q = x2 * q + BETA_0;
77
78    p / q
79}
80
81routine_ew_rust!(generic;
82    f32,
83    generic_tanh_f32_4n,
84    4,
85    4,
86    fn run(x: &mut [f32], _: ()) {
87        debug_assert!(x.len() % Self::nr() == 0);
88        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
89        x.iter_mut().for_each(|px| *px = stanh(*px))
90    },
91    func(Tanh)
92);
93
94routine_ew_rust!(generic;
95    f16,
96    generic_tanh_f16_8n,
97    8,
98    8,
99    fn run(x: &mut [f16], _: ()) {
100        debug_assert!(x.len() % Self::nr() == 0);
101        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
102        x.iter_mut().for_each(|px| *px = htanh(*px))
103    },
104    func(Tanh)
105);