Skip to main content

tract_linalg/generic/
tanh.rs

1#![allow(clippy::excessive_precision)]
2use crate::frame::element_wise::ElementWiseKer;
3use tract_data::internal::*;
4
5/// f32 tanh, as a rational minimax fit of `tanh` over `[LOW, HIGH]`.
6///
7/// The quotient is clamped to `[-1, 1]` to keep the range consumers rely on. Across the
8/// top of the input range the correctly rounded result is already within one ulp of `±1`
9/// — `1 - tanh(8.9)` is `3.7e-8`, under the `2^-24` f32 step just below 1 — while the two
10/// Horner chains and the division leave a few ulps of relative error, so an upward
11/// rounding lands on `±(1 + 2^-22)`. NaN still propagates.
12pub fn stanh(x: f32) -> f32 {
13    const LOW: f32 = -8.9;
14    const HIGH: f32 = 8.9;
15
16    const ALPHA_13: f32 = -8.488492677e-14;
17    const ALPHA_11: f32 = 5.277853000e-11;
18    const ALPHA_9: f32 = -2.022500419e-8;
19    const ALPHA_7: f32 = 0.00001115424833;
20    const ALPHA_5: f32 = 0.003103950131;
21    const ALPHA_3: f32 = 0.1308400453;
22    const ALPHA_1: f32 = 0.9999999934;
23
24    const BETA_6: f32 = 0.0002546136580;
25    const BETA_4: f32 = 0.02449515379;
26    const BETA_2: f32 = 0.4641733162;
27    const BETA_0: f32 = 1.0;
28
29    let x = x.clamp(LOW, HIGH);
30
31    let x2 = x * x;
32
33    let p = ALPHA_13;
34    let p = x2 * p + ALPHA_11;
35    let p = x2 * p + ALPHA_9;
36    let p = x2 * p + ALPHA_7;
37    let p = x2 * p + ALPHA_5;
38    let p = x2 * p + ALPHA_3;
39    let p = x2 * p + ALPHA_1;
40    let p = p * x;
41
42    let q = BETA_6;
43    let q = x2 * q + BETA_4;
44    let q = x2 * q + BETA_2;
45    let q = x2 * q + BETA_0;
46
47    (p / q).clamp(-1.0, 1.0)
48}
49
50/// f16 tanh, as a rational minimax fit of `tanh` over `[LOW, HIGH]`.
51///
52/// Needs no output clamp, unlike [`stanh`]: `1 - tanh(3.84)` is `9.3e-4`, about two f16
53/// steps below 1, so the quotient keeps its margin for every f16 input.
54pub fn htanh(x: f16) -> f16 {
55    const LOW: f16 = f16::from_f32_const(-3.84);
56    const HIGH: f16 = f16::from_f32_const(3.84);
57
58    const ALPHA_3: f16 = f16::from_f32_const(0.082654955);
59    const ALPHA_1: f16 = f16::from_f32_const(0.99963124);
60
61    const BETA_4: f16 = f16::from_f32_const(0.0065383179);
62    const BETA_2: f16 = f16::from_f32_const(0.41401828);
63    const BETA_0: f16 = f16::from_f32_const(1.0);
64
65    let x = x.clamp(LOW, HIGH);
66
67    let x2 = x * x;
68
69    let p = ALPHA_3;
70    let p = x2 * p + ALPHA_1;
71    let p = p * x;
72
73    let q = BETA_4;
74    let q = x2 * q + BETA_2;
75    let q = x2 * q + BETA_0;
76
77    p / q
78}
79
80#[derive(Clone, Debug)]
81pub struct STanh4;
82
83impl ElementWiseKer<f32> for STanh4 {
84    fn name() -> &'static str {
85        "generic"
86    }
87
88    fn alignment_items() -> usize {
89        16
90    }
91
92    fn alignment_bytes() -> usize {
93        16
94    }
95
96    fn nr() -> usize {
97        4
98    }
99
100    fn run(x: &mut [f32], _: ()) {
101        debug_assert!(x.len() % Self::nr() == 0);
102        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
103        x.iter_mut().for_each(|px| *px = stanh(*px))
104    }
105}
106
107#[cfg(test)]
108#[macro_use]
109pub mod s {
110    tanh_frame_tests!(true, f32, crate::generic::tanh::STanh4);
111}
112
113#[derive(Clone, Debug)]
114pub struct HTanh8;
115
116impl ElementWiseKer<f16> for HTanh8 {
117    fn name() -> &'static str {
118        "generic"
119    }
120
121    fn alignment_items() -> usize {
122        16
123    }
124
125    fn alignment_bytes() -> usize {
126        16
127    }
128
129    fn nr() -> usize {
130        8
131    }
132
133    fn run(x: &mut [f16], _: ()) {
134        debug_assert!(x.len() % Self::nr() == 0);
135        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
136        x.iter_mut().for_each(|px| *px = htanh(*px))
137    }
138}
139
140#[cfg(test)]
141#[macro_use]
142pub mod h {
143    tanh_frame_tests!(true, tract_data::internal::f16, crate::generic::tanh::HTanh8);
144}