ferric_tensor/
autograd.rs

1//! Reverse-mode automatic differentiation over the general tensor runtime — the layer that turns
2//! Ferric from an inference demo into a fabric that can **train**. A `Var` wraps a `Tensor` and
3//! records how to backpropagate through each op; `backward()` walks the graph in reverse and
4//! accumulates gradients (broadcasting-aware). Params live as plain `Tensor`s and are re-wrapped
5//! each step, so an optimizer is just tensor arithmetic.
6//!
7//! Validated by a finite-difference gradient check and by actually training an MLP (loss ↓).
8
9use crate::Tensor;
10use ferric_core::Context;
11use std::cell::RefCell;
12use std::rc::Rc;
13use std::sync::Arc;
14
15type BackFn = dyn Fn(&Tensor, &[Var]);
16
17struct Inner {
18    value: Tensor,
19    grad: RefCell<Option<Tensor>>,
20    parents: Vec<Var>,
21    backward: Option<Box<BackFn>>,
22}
23
24#[derive(Clone)]
25pub struct Var(Rc<Inner>);
26
27impl Var {
28    /// A differentiable leaf (a parameter or input we want gradients for).
29    pub fn leaf(t: Tensor) -> Var {
30        Var(Rc::new(Inner { value: t, grad: RefCell::new(None), parents: vec![], backward: None }))
31    }
32    fn node(value: Tensor, parents: Vec<Var>, backward: Box<BackFn>) -> Var {
33        Var(Rc::new(Inner { value, grad: RefCell::new(None), parents, backward: Some(backward) }))
34    }
35    pub fn value(&self) -> &Tensor { &self.0.value }
36    pub fn grad(&self) -> Option<Tensor> { self.0.grad.borrow().clone() }
37
38    fn ctx(&self) -> Arc<Context> { self.0.value.ctx_arc() }
39    fn accumulate(&self, g: &Tensor) {
40        let g = unbroadcast(g, &self.0.value.shape);
41        let mut slot = self.0.grad.borrow_mut();
42        *slot = Some(match slot.take() {
43            Some(prev) => prev.add(&g),
44            None => g,
45        });
46    }
47
48    /// Backpropagate from this (scalar-ish) output: seed grad = ones, walk reverse-topo.
49    pub fn backward(&self) {
50        // post-order DFS → topo order
51        let mut topo: Vec<Var> = vec![];
52        let mut seen: Vec<*const Inner> = vec![];
53        fn visit(v: &Var, topo: &mut Vec<Var>, seen: &mut Vec<*const Inner>) {
54            let p = Rc::as_ptr(&v.0);
55            if seen.contains(&p) { return; }
56            seen.push(p);
57            for par in &v.0.parents { visit(par, topo, seen); }
58            topo.push(v.clone());
59        }
60        visit(self, &mut topo, &mut seen);
61        let ones = Tensor::from_vec(&self.ctx(), &vec![1.0; self.0.value.numel()], &self.0.value.shape);
62        *self.0.grad.borrow_mut() = Some(ones);
63        for v in topo.iter().rev() {
64            if let Some(bw) = &v.0.backward {
65                let g = v.0.grad.borrow().clone().expect("grad set before backward");
66                bw(&g, &v.0.parents);
67            }
68        }
69    }
70
71    // ---- differentiable ops ----
72    pub fn add(&self, o: &Var) -> Var {
73        let out = self.0.value.add(&o.0.value);
74        Var::node(out, vec![self.clone(), o.clone()], Box::new(|g, p| { p[0].accumulate(g); p[1].accumulate(g); }))
75    }
76    pub fn sub(&self, o: &Var) -> Var {
77        let out = self.0.value.sub(&o.0.value);
78        Var::node(out, vec![self.clone(), o.clone()], Box::new(|g, p| { p[0].accumulate(g); p[1].accumulate(&g.neg()); }))
79    }
80    pub fn mul(&self, o: &Var) -> Var {
81        let out = self.0.value.mul(&o.0.value);
82        let (a, b) = (self.0.value.clone(), o.0.value.clone());
83        Var::node(out, vec![self.clone(), o.clone()], Box::new(move |g, p| {
84            p[0].accumulate(&g.mul(&b));
85            p[1].accumulate(&g.mul(&a));
86        }))
87    }
88    /// Matmul (last two dims; equal or absent batch dims).
89    pub fn matmul(&self, o: &Var) -> Var {
90        let out = self.0.value.matmul(&o.0.value);
91        let (a, b) = (self.0.value.clone(), o.0.value.clone());
92        Var::node(out, vec![self.clone(), o.clone()], Box::new(move |g, p| {
93            let r = a.rank();
94            p[0].accumulate(&g.matmul(&b.transpose(r - 1, r - 2)));   // dA = g · Bᵀ
95            p[1].accumulate(&a.transpose(r - 1, r - 2).matmul(g));    // dB = Aᵀ · g
96        }))
97    }
98    pub fn relu(&self) -> Var {
99        let out = self.0.value.relu();
100        let x = self.0.value.clone();
101        Var::node(out, vec![self.clone()], Box::new(move |g, p| {
102            // grad * (x > 0)  ==  grad * (relu(x)/x) avoided; use step = (x>0) via max(sign,0)
103            let mask = x.relu_mask();
104            p[0].accumulate(&g.mul(&mask));
105        }))
106    }
107    pub fn sum_all(&self) -> Var {
108        let axes: Vec<usize> = (0..self.0.value.rank()).collect();
109        let out = self.0.value.sum(&axes, false);
110        let shape = self.0.value.shape.clone();
111        Var::node(out, vec![self.clone()], Box::new(move |g, p| {
112            // g is scalar [1]; broadcast to input shape
113            let full = Tensor::from_vec(&p[0].ctx(), &vec![0.0; numel(&shape)], &shape);
114            p[0].accumulate(&full.add(g)); // add broadcasts scalar over full
115        }))
116    }
117    pub fn mean_all(&self) -> Var {
118        let n = self.0.value.numel() as f32;
119        let s = self.sum_all();
120        let inv = Var::leaf(Tensor::from_vec(&self.ctx(), &[1.0 / n], &[1]));
121        s.mul(&inv)
122    }
123}
124
125fn numel(s: &[usize]) -> usize { s.iter().product() }
126
127/// Sum a gradient back down to `shape` (reverse of broadcasting).
128fn unbroadcast(g: &Tensor, shape: &[usize]) -> Tensor {
129    let mut out = g.clone();
130    // collapse extra leading dims
131    while out.rank() > shape.len() {
132        out = out.sum(&[0], false);
133    }
134    // sum dims that were broadcast (target size 1, grad size > 1)
135    let axes: Vec<usize> = (0..shape.len()).filter(|&d| shape[d] == 1 && out.shape[d] != 1).collect();
136    if !axes.is_empty() {
137        out = out.sum(&axes, true);
138    }
139    if out.shape != shape {
140        out = out.reshape(shape);
141    }
142    out
143}
ferric_tensor/autograd.rs

ferric_tensor/
autograd.rs
