1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
//! `silu` and `gelu` activations, plus the `tanh` and `sigmoid`/`logistic`
//! pair the LSTM OCR recognizer's gates are built from.
use kopitiam_core::{DType, Result};
use crate::storage::Storage;
use super::Tensor;
/// `sqrt(2 / pi)`, the constant in the tanh approximation of GELU.
const SQRT_2_OVER_PI: f32 = 0.797_884_6;
impl Tensor {
/// SiLU / swish (Elfwing, Uchibe & Doya, 2017; also Hendrycks & Gimpel):
/// `x * sigmoid(x) = x / (1 + exp(-x))`. Used by LLaMA-family FFNs.
pub fn silu(&self) -> Result<Tensor> {
self.unary_f32(|x| x / (1.0 + (-x).exp()))
}
/// GELU (Hendrycks & Gimpel, 2016), tanh approximation:
/// `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`.
///
/// This crate implements the tanh approximation rather than the exact
/// `0.5 * x * (1 + erf(x / sqrt(2)))` form because `erf` is not in
/// Rust's `std` — computing it exactly would mean shipping a rational
/// or Chebyshev `erf` approximation of our own, which is *more* code
/// and *more* numerical risk than the tanh form for a difference that
/// is below `1e-3` everywhere and is what GPT-2/GPT-NeoX-family models
/// were themselves trained and evaluated with (the tanh form is
/// sometimes called "gelu_new" in that lineage). If a model that
/// specifically requires the exact erf form shows up, add it as a
/// second method rather than replacing this one.
pub fn gelu(&self) -> Result<Tensor> {
self.unary_f32(|x| 0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044_715 * x * x * x)).tanh()))
}
/// Hyperbolic tangent, elementwise: `tanh(x)`, in `(-1, 1)`. The
/// cell-candidate (`g`) nonlinearity and the output squashing (`tanh(c')`)
/// of an LSTM cell — see [`Tensor::lstm_cell`].
///
/// Matches Tesseract's `Tanh` (`functions.cpp`), computed directly via
/// `f32::tanh` rather than through that file's 8k-entry lookup table
/// (`TanhTable`): standard math, so there is nothing Tesseract-specific to
/// port here — the LUT is a speed/accuracy trade this crate has no reason
/// to reproduce.
pub fn tanh(&self) -> Result<Tensor> {
self.unary_f32(f32::tanh)
}
/// Logistic sigmoid, elementwise: `1 / (1 + exp(-x))`, in `(0, 1)`. The
/// input/forget/output gate nonlinearity of an LSTM cell — see
/// [`Tensor::lstm_cell`].
///
/// Matches Tesseract's `Logistic` (`functions.cpp`), computed directly
/// rather than via its 8k-entry lookup table (`LogisticTable`); standard
/// math, no Tesseract-specific code to port. [`Tensor::logistic`] is an
/// alias for callers that prefer Tesseract's name for the same function.
pub fn sigmoid(&self) -> Result<Tensor> {
self.unary_f32(|x| 1.0 / (1.0 + (-x).exp()))
}
/// Alias for [`Tensor::sigmoid`] under Tesseract's name for the function.
pub fn logistic(&self) -> Result<Tensor> {
self.sigmoid()
}
fn unary_f32(&self, f: impl Fn(f32) -> f32) -> Result<Tensor> {
self.require_dtype(DType::F32)?;
let Storage::F32(data) = self.storage.as_ref() else { unreachable!() };
let out: Vec<f32> = self.logical_offsets().map(|i| f(data[i])).collect();
Tensor::from_f32(out, self.shape.clone())
}
}
#[cfg(test)]
mod tests {
use kopitiam_core::Error;
use super::*;
fn assert_close(a: f32, b: f32, epsilon: f32) {
assert!((a - b).abs() < epsilon, "expected {b}, got {a}");
}
#[test]
fn silu_matches_hand_computation() {
// silu(0) = 0 * sigmoid(0) = 0 * 0.5 = 0
let t = Tensor::from_f32(vec![0.0, 1.0, -1.0], [3]).unwrap();
let out = t.silu().unwrap().to_vec_f32().unwrap();
assert_close(out[0], 0.0, 1e-6);
// silu(1) = 1 / (1 + e^-1) ~= 0.7310586
assert_close(out[1], 0.731_058_6, 1e-5);
// silu(-1) = -1 / (1 + e^1) ~= -0.2689414
assert_close(out[2], -0.268_941_4, 1e-5);
}
#[test]
fn silu_matches_the_reference_formula_across_a_range_of_inputs() {
// silu is *not* globally monotonic (it dips to a minimum around
// x ~= -1.278 before rising) so the meaningful correctness check is
// "matches x * sigmoid(x) pointwise", not a monotonicity claim.
let xs = vec![-5.0, -3.0, -1.278, -1.0, -0.5, 0.0, 0.5, 1.0, 3.0, 5.0];
let t = Tensor::from_f32(xs.clone(), [xs.len()]).unwrap();
let out = t.silu().unwrap().to_vec_f32().unwrap();
for (x, o) in xs.iter().zip(&out) {
let expected = x / (1.0 + (-x).exp());
assert_close(*o, expected, 1e-5);
}
}
#[test]
fn silu_is_monotonically_increasing_on_its_increasing_branch() {
// For x >= -1.278ish, silu is monotonically increasing; this is the
// range every real activation input in a trained model's residual
// stream overwhelmingly falls into after the first few layers.
let t = Tensor::from_f32(vec![-1.0, 0.0, 1.0, 2.0, 3.0], [5]).unwrap();
let out = t.silu().unwrap().to_vec_f32().unwrap();
for pair in out.windows(2) {
assert!(pair[1] > pair[0], "silu should be increasing here: {out:?}");
}
}
#[test]
fn gelu_matches_hand_computation_at_zero_and_matches_known_values() {
// gelu(0) = 0.
let t = Tensor::from_f32(vec![0.0, 1.0, -1.0], [3]).unwrap();
let out = t.gelu().unwrap().to_vec_f32().unwrap();
assert_close(out[0], 0.0, 1e-6);
// Known reference values for the tanh-approximation GELU (matches
// the widely used "gelu_new" implementation to ~1e-6).
assert_close(out[1], 0.841_192, 1e-4);
assert_close(out[2], -0.158_808, 1e-4);
}
#[test]
fn gelu_approaches_the_identity_for_large_positive_x_and_zero_for_large_negative_x() {
let t = Tensor::from_f32(vec![10.0, -10.0], [2]).unwrap();
let out = t.gelu().unwrap().to_vec_f32().unwrap();
assert_close(out[0], 10.0, 1e-3);
assert_close(out[1], 0.0, 1e-3);
}
#[test]
fn activations_reject_non_f32_input() {
let t = Tensor::from_i32(vec![1, 2, 3], [3]).unwrap();
assert!(matches!(t.silu(), Err(Error::DTypeMismatch { .. })));
assert!(matches!(t.gelu(), Err(Error::DTypeMismatch { .. })));
assert!(matches!(t.tanh(), Err(Error::DTypeMismatch { .. })));
assert!(matches!(t.sigmoid(), Err(Error::DTypeMismatch { .. })));
}
#[test]
fn tanh_matches_known_values_and_is_odd() {
// tanh(0) = 0; tanh(1) ~= 0.7615942; tanh is odd: tanh(-1) = -tanh(1).
let t = Tensor::from_f32(vec![0.0, 1.0, -1.0], [3]).unwrap();
let out = t.tanh().unwrap().to_vec_f32().unwrap();
assert_close(out[0], 0.0, 1e-6);
assert_close(out[1], 0.761_594_2, 1e-6);
assert_close(out[2], -0.761_594_2, 1e-6);
}
#[test]
fn tanh_is_monotonic_and_saturates_towards_plus_minus_one() {
let t = Tensor::from_f32(vec![-20.0, -2.0, -0.5, 0.0, 0.5, 2.0, 20.0], [7]).unwrap();
let out = t.tanh().unwrap().to_vec_f32().unwrap();
for pair in out.windows(2) {
assert!(pair[1] > pair[0], "tanh should be strictly increasing: {out:?}");
}
assert_close(out[0], -1.0, 1e-6); // saturates to -1 for large negative x.
assert_close(out[6], 1.0, 1e-6); // saturates to +1 for large positive x.
}
#[test]
fn sigmoid_matches_known_values_and_the_reflection_identity() {
// sigmoid(0) = 0.5; sigmoid(2) ~= 0.8807971; sigmoid(-x) = 1 - sigmoid(x).
let t = Tensor::from_f32(vec![0.0, 2.0, -2.0], [3]).unwrap();
let out = t.sigmoid().unwrap().to_vec_f32().unwrap();
assert_close(out[0], 0.5, 1e-6);
assert_close(out[1], 0.880_797_1, 1e-6);
assert_close(out[2], 1.0 - 0.880_797_1, 1e-6);
}
#[test]
fn sigmoid_is_monotonic_saturates_and_stays_in_the_open_unit_interval() {
let t = Tensor::from_f32(vec![-40.0, -3.0, 0.0, 3.0, 40.0], [5]).unwrap();
let out = t.sigmoid().unwrap().to_vec_f32().unwrap();
for pair in out.windows(2) {
assert!(pair[1] > pair[0], "sigmoid should be strictly increasing: {out:?}");
}
for &v in &out {
assert!((0.0..=1.0).contains(&v), "sigmoid output escaped [0, 1]: {v}");
}
assert_close(out[0], 0.0, 1e-6); // saturates towards 0.
assert_close(out[4], 1.0, 1e-6); // saturates towards 1.
}
#[test]
fn logistic_is_an_alias_for_sigmoid() {
let t = Tensor::from_f32(vec![-1.5, 0.0, 0.25, 3.0], [4]).unwrap();
assert_eq!(
t.logistic().unwrap().to_vec_f32().unwrap(),
t.sigmoid().unwrap().to_vec_f32().unwrap(),
);
}
#[test]
fn tanh_and_sigmoid_preserve_shape_and_respect_a_transposed_view() {
// Elementwise ops must read through logical order, not raw storage.
let t = Tensor::from_f32(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0], [2, 3]).unwrap();
let tt = t.transpose(0, 1).unwrap(); // shape [3, 2], non-contiguous.
let out = tt.tanh().unwrap();
assert_eq!(out.shape().dims(), &[3, 2]);
assert_eq!(
out.to_vec_f32().unwrap(),
[0.0f32, 3.0, 1.0, 4.0, 2.0, 5.0].iter().map(|x| x.tanh()).collect::<Vec<_>>(),
);
assert_eq!(tt.sigmoid().unwrap().shape().dims(), &[3, 2]);
}
}