teeny_kernels/nn/activation/
elu.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// ── ELU ──────────────────────────────────────────────────────────────────────
27
28/// Forward: y = x if x > 0 else alpha*(exp(x) - 1)
29#[kernel(backward = EluBackward)]
30pub fn elu_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    alpha: f32,
35) where
36    T::I32Tensor: types::Tensor<i32, 1>,
37    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
38    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
39{
40    let pid = T::program_id(Axis::X);
41    let block_start = pid * BLOCK_SIZE;
42    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
43    let in_bounds = offsets.lt(n_elements);
44
45    let x = T::load(
46        x_ptr.add_offsets(offsets),
47        Some(in_bounds),
48        None,
49        &[],
50        None,
51        None,
52        None,
53        false,
54    );
55    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
56    let alpha_t = T::full(&[BLOCK_SIZE], D::from_f64(alpha as f64));
57    let x_pos = T::gt(x, T::zeros_like(x));
58    let y = T::where_(x_pos, x, alpha_t * (T::exp(x) - one));
59    T::store(
60        y_ptr.add_offsets(offsets),
61        y,
62        Some(in_bounds),
63        &[],
64        None,
65        None,
66    );
67}
68
69/// Backward: dx = dy if x > 0 else dy * alpha * exp(x)
70#[kernel]
71pub fn elu_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
72    dy_ptr: T::Pointer<D>,
73    x_ptr: T::Pointer<D>,
74    dx_ptr: T::Pointer<D>,
75    n_elements: i32,
76    alpha: f32,
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 x = T::load(
98        x_ptr.add_offsets(offsets),
99        Some(in_bounds),
100        None,
101        &[],
102        None,
103        None,
104        None,
105        false,
106    );
107    let alpha_t = T::full(&[BLOCK_SIZE], D::from_f64(alpha as f64));
108    let x_pos = T::gt(x, T::zeros_like(x));
109    let dx = T::where_(x_pos, dy, dy * alpha_t * T::exp(x));
110    T::store(
111        dx_ptr.add_offsets(offsets),
112        dx,
113        Some(in_bounds),
114        &[],
115        None,
116        None,
117    );
118}
119
120// ── SELU ─────────────────────────────────────────────────────────────────────
121
122/// Forward: y = SCALE * (x if x > 0 else ALPHA*(exp(x) - 1))
123#[kernel(backward = SeluBackward)]
124pub fn selu_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
125    x_ptr: T::Pointer<D>,
126    y_ptr: T::Pointer<D>,
127    n_elements: i32,
128) where
129    T::I32Tensor: types::Tensor<i32, 1>,
130    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
131    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
132{
133    let pid = T::program_id(Axis::X);
134    let block_start = pid * BLOCK_SIZE;
135    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
136    let in_bounds = offsets.lt(n_elements);
137
138    let x = T::load(
139        x_ptr.add_offsets(offsets),
140        Some(in_bounds),
141        None,
142        &[],
143        None,
144        None,
145        None,
146        false,
147    );
148    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
149    let scale = T::full(&[BLOCK_SIZE], D::from_f64(1.0507009873554804));
150    let alpha = T::full(&[BLOCK_SIZE], D::from_f64(1.6732632423543772));
151    let x_pos = T::gt(x, T::zeros_like(x));
152    let y = scale * T::where_(x_pos, x, alpha * (T::exp(x) - one));
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 = SCALE*dy if x > 0 else dy * SCALE*ALPHA*exp(x)
164#[kernel]
165pub fn selu_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
166    dy_ptr: T::Pointer<D>,
167    x_ptr: T::Pointer<D>,
168    dx_ptr: T::Pointer<D>,
169    n_elements: i32,
170) where
171    T::I32Tensor: types::Tensor<i32, 1>,
172    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
173    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
174{
175    let pid = T::program_id(Axis::X);
176    let block_start = pid * BLOCK_SIZE;
177    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
178    let in_bounds = offsets.lt(n_elements);
179
180    let dy = T::load(
181        dy_ptr.add_offsets(offsets),
182        Some(in_bounds),
183        None,
184        &[],
185        None,
186        None,
187        None,
188        false,
189    );
190    let x = T::load(
191        x_ptr.add_offsets(offsets),
192        Some(in_bounds),
193        None,
194        &[],
195        None,
196        None,
197        None,
198        false,
199    );
200    let scale = T::full(&[BLOCK_SIZE], D::from_f64(1.0507009873554804));
201    let scale_alpha = T::full(&[BLOCK_SIZE], D::from_f64(1.7580993408473766));
202    let x_pos = T::gt(x, T::zeros_like(x));
203    let dx = T::where_(x_pos, dy * scale, dy * scale_alpha * T::exp(x));
204    T::store(
205        dx_ptr.add_offsets(offsets),
206        dx,
207        Some(in_bounds),
208        &[],
209        None,
210        None,
211    );
212}
213
214// ── CELU ─────────────────────────────────────────────────────────────────────
215
216/// Forward: y = max(0, x) + min(0, alpha*(exp(x/alpha) - 1))
217#[kernel(backward = CeluBackward)]
218pub fn celu_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
219    x_ptr: T::Pointer<D>,
220    y_ptr: T::Pointer<D>,
221    n_elements: i32,
222    alpha: f32,
223) where
224    T::I32Tensor: types::Tensor<i32, 1>,
225    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
226    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
227{
228    let pid = T::program_id(Axis::X);
229    let block_start = pid * BLOCK_SIZE;
230    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
231    let in_bounds = offsets.lt(n_elements);
232
233    let x = T::load(
234        x_ptr.add_offsets(offsets),
235        Some(in_bounds),
236        None,
237        &[],
238        None,
239        None,
240        None,
241        false,
242    );
243    let zero = T::zeros_like(x);
244    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
245    let alpha_t = T::full(&[BLOCK_SIZE], D::from_f64(alpha as f64));
246    let inv_alpha = T::full(&[BLOCK_SIZE], D::from_f64(1.0 / alpha as f64));
247    let elu_neg = alpha_t * (T::exp(x * inv_alpha) - one);
248    let y = T::maximum(zero, x) + T::minimum(zero, elu_neg);
249    T::store(
250        y_ptr.add_offsets(offsets),
251        y,
252        Some(in_bounds),
253        &[],
254        None,
255        None,
256    );
257}
258
259/// Backward: dx = dy if x >= 0 else dy * exp(x/alpha)
260#[kernel]
261pub fn celu_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
262    dy_ptr: T::Pointer<D>,
263    x_ptr: T::Pointer<D>,
264    dx_ptr: T::Pointer<D>,
265    n_elements: i32,
266    alpha: f32,
267) where
268    T::I32Tensor: types::Tensor<i32, 1>,
269    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
270    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
271{
272    let pid = T::program_id(Axis::X);
273    let block_start = pid * BLOCK_SIZE;
274    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
275    let in_bounds = offsets.lt(n_elements);
276
277    let dy = T::load(
278        dy_ptr.add_offsets(offsets),
279        Some(in_bounds),
280        None,
281        &[],
282        None,
283        None,
284        None,
285        false,
286    );
287    let x = T::load(
288        x_ptr.add_offsets(offsets),
289        Some(in_bounds),
290        None,
291        &[],
292        None,
293        None,
294        None,
295        false,
296    );
297    let inv_alpha = T::full(&[BLOCK_SIZE], D::from_f64(1.0 / alpha as f64));
298    let x_ge_zero = T::ge(x, T::zeros_like(x));
299    let dx = T::where_(x_ge_zero, dy, dy * T::exp(x * inv_alpha));
300    T::store(
301        dx_ptr.add_offsets(offsets),
302        dx,
303        Some(in_bounds),
304        &[],
305        None,
306        None,
307    );
308}
309
310pub struct EluOp<D: Float> {
311    pub forward: EluForward<D>,
312    pub backward: EluBackward<D>,
313}
314
315pub struct SeluOp<D: Float> {
316    pub forward: SeluForward<D>,
317    pub backward: SeluBackward<D>,
318}
319
320pub struct CeluOp<D: Float> {
321    pub forward: CeluForward<D>,
322    pub backward: CeluBackward<D>,
323}
teeny_kernels/nn/activation/elu.rs

teeny_kernels/nn/activation/
elu.rs
