teeny_kernels/nn/optim/
adamax.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_macros::kernel;
20use teeny_triton::triton::{
21    types::{AddOffsets, Comparison},
22    *,
23};
24
25/// Adamax step (infinity-norm variant of Adam).
26///
27/// ```text
28/// exp_avg = beta1 * exp_avg + (1 - beta1) * g
29/// exp_inf = max(beta2 * exp_inf, |g| + eps)
30/// p       = p - clr * exp_avg / exp_inf
31/// ```
32///
33/// `clr = lr / (1 - beta1^t)` is precomputed on the host.
34///
35/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
36#[kernel]
37pub fn adamax_step<T: Triton, const BLOCK_SIZE: i32>(
38    params_ptr: T::Pointer<f32>,
39    grad_ptr: T::Pointer<f32>,
40    exp_avg_ptr: T::Pointer<f32>,
41    exp_inf_ptr: T::Pointer<f32>,
42    n_elements: i32,
43    clr: f32,
44    beta1: f32,
45    beta2: f32,
46    eps: f32,
47    weight_decay: f32,
48) where
49    T::I32Tensor: types::Tensor<i32, 1>,
50    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
51    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
52{
53    let pid = T::program_id(Axis::X);
54    let block_start = pid * BLOCK_SIZE;
55    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
56    let mask = offsets.lt(n_elements);
57
58    let p = T::load(
59        params_ptr.add_offsets(offsets),
60        Some(mask),
61        None,
62        &[],
63        None,
64        None,
65        None,
66        false,
67    );
68    let g = T::load(
69        grad_ptr.add_offsets(offsets),
70        Some(mask),
71        None,
72        &[],
73        None,
74        None,
75        None,
76        false,
77    );
78    let exp_avg = T::load(
79        exp_avg_ptr.add_offsets(offsets),
80        Some(mask),
81        None,
82        &[],
83        None,
84        None,
85        None,
86        false,
87    );
88    let exp_inf = T::load(
89        exp_inf_ptr.add_offsets(offsets),
90        Some(mask),
91        None,
92        &[],
93        None,
94        None,
95        None,
96        false,
97    );
98
99    let beta1_t = T::full(&[BLOCK_SIZE], beta1);
100    let one_m_beta1 = T::full(&[BLOCK_SIZE], 1.0_f32 - beta1);
101    let beta2_t = T::full(&[BLOCK_SIZE], beta2);
102    let eps_t = T::full(&[BLOCK_SIZE], eps);
103    let clr_t = T::full(&[BLOCK_SIZE], clr);
104    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
105
106    let g_eff = g + wd_t * p;
107    let exp_avg_new = beta1_t * exp_avg + one_m_beta1 * g_eff;
108    // exp_inf = max(beta2 * exp_inf, |g| + eps)
109    let exp_inf_new = T::maximum(beta2_t * exp_inf, T::abs(g_eff) + eps_t);
110    let p_new = p - clr_t * exp_avg_new / exp_inf_new;
111
112    T::store(
113        params_ptr.add_offsets(offsets),
114        p_new,
115        Some(mask),
116        &[],
117        None,
118        None,
119    );
120    T::store(
121        exp_avg_ptr.add_offsets(offsets),
122        exp_avg_new,
123        Some(mask),
124        &[],
125        None,
126        None,
127    );
128    T::store(
129        exp_inf_ptr.add_offsets(offsets),
130        exp_inf_new,
131        Some(mask),
132        &[],
133        None,
134        None,
135    );
136}
teeny_kernels/nn/optim/adamax.rs

teeny_kernels/nn/optim/
adamax.rs
