teeny_kernels/nn/activation/
tanh.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#![allow(non_snake_case)]
18
19use teeny_core::dtype::Float;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison},
23    *,
24};
25
26// ── Tanh ─────────────────────────────────────────────────────────────────────
27
28/// Forward: y = tanh(x) = 2*sigmoid(2x) - 1 = 2/(1+exp(-2x)) - 1
29#[kernel(backward = TanhBackward)]
30pub fn tanh_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
31    x_ptr: T::Pointer<D>,
32    y_ptr: T::Pointer<D>,
33    n_elements: i32,
34) where
35    T::I32Tensor: types::Tensor<i32, 1>,
36    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
37    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
38{
39    let pid = T::program_id(Axis::X);
40    let block_start = pid * BLOCK_SIZE;
41    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
42    let in_bounds = offsets.lt(n_elements);
43
44    let x = T::load(
45        x_ptr.add_offsets(offsets),
46        Some(in_bounds),
47        None,
48        &[],
49        None,
50        None,
51        None,
52        false,
53    );
54    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
55    let two = T::full(&[BLOCK_SIZE], D::from_f64(2.0));
56    let neg2 = T::full(&[BLOCK_SIZE], D::from_f64(-2.0));
57    // sigmoid(2x) = 1 / (1 + exp(-2x))
58    let s2x = one / (one + T::exp(neg2 * x));
59    let y = two * s2x - one;
60    T::store(
61        y_ptr.add_offsets(offsets),
62        y,
63        Some(in_bounds),
64        &[],
65        None,
66        None,
67    );
68}
69
70/// Backward: dx = dy * (1 - y²)  — sech²(x) expressed via saved output
71#[kernel]
72pub fn tanh_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
73    dy_ptr: T::Pointer<D>,
74    y_ptr: T::Pointer<D>,
75    dx_ptr: T::Pointer<D>,
76    n_elements: i32,
77) where
78    T::I32Tensor: types::Tensor<i32, 1>,
79    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
80    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
81{
82    let pid = T::program_id(Axis::X);
83    let block_start = pid * BLOCK_SIZE;
84    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
85    let in_bounds = offsets.lt(n_elements);
86
87    let dy = T::load(
88        dy_ptr.add_offsets(offsets),
89        Some(in_bounds),
90        None,
91        &[],
92        None,
93        None,
94        None,
95        false,
96    );
97    let y = T::load(
98        y_ptr.add_offsets(offsets),
99        Some(in_bounds),
100        None,
101        &[],
102        None,
103        None,
104        None,
105        false,
106    );
107    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
108    let dx = dy * (one - y * y);
109    T::store(
110        dx_ptr.add_offsets(offsets),
111        dx,
112        Some(in_bounds),
113        &[],
114        None,
115        None,
116    );
117}
118
119// ── Tanhshrink ───────────────────────────────────────────────────────────────
120
121/// Forward: y = x - tanh(x)
122#[kernel(backward = TanhshrinkBackward)]
123pub fn tanhshrink_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
124    x_ptr: T::Pointer<D>,
125    y_ptr: T::Pointer<D>,
126    n_elements: i32,
127) where
128    T::I32Tensor: types::Tensor<i32, 1>,
129    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
130    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
131{
132    let pid = T::program_id(Axis::X);
133    let block_start = pid * BLOCK_SIZE;
134    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
135    let in_bounds = offsets.lt(n_elements);
136
137    let x = T::load(
138        x_ptr.add_offsets(offsets),
139        Some(in_bounds),
140        None,
141        &[],
142        None,
143        None,
144        None,
145        false,
146    );
147    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
148    let two = T::full(&[BLOCK_SIZE], D::from_f64(2.0));
149    let neg2 = T::full(&[BLOCK_SIZE], D::from_f64(-2.0));
150    let s2x = one / (one + T::exp(neg2 * x));
151    let tanh_x = two * s2x - one;
152    let y = x - tanh_x;
153    T::store(
154        y_ptr.add_offsets(offsets),
155        y,
156        Some(in_bounds),
157        &[],
158        None,
159        None,
160    );
161}
162
163/// Backward: dx = dy * tanh²(x)
164///   Since y = x - tanh(x), we have tanh(x) = x - y, so tanh²(x) = (x-y)².
165#[kernel]
166pub fn tanhshrink_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
167    dy_ptr: T::Pointer<D>,
168    x_ptr: T::Pointer<D>,
169    y_ptr: T::Pointer<D>,
170    dx_ptr: T::Pointer<D>,
171    n_elements: i32,
172) where
173    T::I32Tensor: types::Tensor<i32, 1>,
174    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
175    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
176{
177    let pid = T::program_id(Axis::X);
178    let block_start = pid * BLOCK_SIZE;
179    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
180    let in_bounds = offsets.lt(n_elements);
181
182    let dy = T::load(
183        dy_ptr.add_offsets(offsets),
184        Some(in_bounds),
185        None,
186        &[],
187        None,
188        None,
189        None,
190        false,
191    );
192    let x = T::load(
193        x_ptr.add_offsets(offsets),
194        Some(in_bounds),
195        None,
196        &[],
197        None,
198        None,
199        None,
200        false,
201    );
202    let y = T::load(
203        y_ptr.add_offsets(offsets),
204        Some(in_bounds),
205        None,
206        &[],
207        None,
208        None,
209        None,
210        false,
211    );
212    let tanh_x = x - y;
213    let dx = dy * tanh_x * tanh_x;
214    T::store(
215        dx_ptr.add_offsets(offsets),
216        dx,
217        Some(in_bounds),
218        &[],
219        None,
220        None,
221    );
222}
223
224pub struct TanhOp<D: Float> {
225    pub forward: TanhForward<D>,
226    pub backward: TanhBackward<D>,
227}
228
229pub struct TanhshrinkOp<D: Float> {
230    pub forward: TanhshrinkForward<D>,
231    pub backward: TanhshrinkBackward<D>,
232}
teeny_kernels/nn/activation/tanh.rs

teeny_kernels/nn/activation/
tanh.rs
