ferric_tensor/
optim.rs

1//! Optimizers — plain tensor arithmetic over the general runtime. Adam keeps per-parameter first/
2//! second moment estimates and applies bias-corrected updates entirely on the GPU.
3
4use crate::Tensor;
5
6pub struct Adam {
7    lr: f32,
8    b1: f32,
9    b2: f32,
10    eps: f32,
11    t: i32,
12    m: Vec<Tensor>,
13    v: Vec<Tensor>,
14}
15
16impl Adam {
17    pub fn new(params: &[Tensor], lr: f32) -> Adam {
18        let m = params.iter().map(|p| Tensor::zeros(&p.ctx_arc(), &p.shape)).collect();
19        let v = params.iter().map(|p| Tensor::zeros(&p.ctx_arc(), &p.shape)).collect();
20        Adam { lr, b1: 0.9, b2: 0.999, eps: 1e-8, t: 0, m, v }
21    }
22
23    /// One update step: `params[i] -= lr · m̂ / (√v̂ + eps)`, replacing each param tensor in place.
24    pub fn step(&mut self, params: &mut [Tensor], grads: &[Tensor]) {
25        self.t += 1;
26        let bc1 = 1.0 / (1.0 - self.b1.powi(self.t));
27        let bc2 = 1.0 / (1.0 - self.b2.powi(self.t));
28        for i in 0..params.len() {
29            let g = &grads[i];
30            let sc = |t: &Tensor, s: f32| t.mul(&t.scalar(s));
31            // m = b1·m + (1-b1)·g ;  v = b2·v + (1-b2)·g²
32            self.m[i] = sc(&self.m[i], self.b1).add(&sc(g, 1.0 - self.b1));
33            self.v[i] = sc(&self.v[i], self.b2).add(&sc(&g.mul(g), 1.0 - self.b2));
34            let mhat = sc(&self.m[i], bc1);
35            let vhat = sc(&self.v[i], bc2);
36            let update = mhat.div(&vhat.sqrt().add(&vhat.scalar(self.eps)));
37            params[i] = params[i].sub(&sc(&update, self.lr));
38        }
39    }
40}
ferric_tensor/optim.rs

ferric_tensor/
optim.rs
