tang_expr/
reverse.rs

1//! Reverse-mode differentiation: vector-Jacobian products as graphs.
2
3use std::collections::HashMap;
4
5use crate::graph::ExprGraph;
6use crate::node::{ExprId, Node};
7
8impl ExprGraph {
9    /// The vector-Jacobian product `seedᵀ ∂outputs/∂vars` as new
10    /// expressions in this graph: entry `k` is `Σ_i seeds[i] ∂outputs[i]/∂Var(k)`
11    /// for `k < n_vars`. One reverse sweep over the live nodes, so the cost
12    /// is a small multiple of evaluating `outputs` whatever `n_vars` is
13    /// (forward [`diff`](Self::diff) is one sweep per variable). Seeds are
14    /// expressions too (typically fresh `Var`s past the inputs), so one
15    /// compiled graph serves every seed. `Select` passes gradient to the
16    /// taken branch only, as `diff` does.
17    pub fn vjp(&mut self, outputs: &[ExprId], seeds: &[ExprId], n_vars: usize) -> Vec<ExprId> {
18        assert_eq!(outputs.len(), seeds.len(), "one seed per output");
19        let live = self.live_set(outputs);
20        let mut order: Vec<usize> = live.into_iter().collect();
21        // children are created before their parents: descending index is a
22        // reverse topological order
23        order.sort_unstable_by(|a, b| b.cmp(a));
24        let mut adj: HashMap<usize, ExprId> = HashMap::new();
25        for (&o, &s) in outputs.iter().zip(seeds) {
26            self.accumulate(&mut adj, o, s);
27        }
28        let mut grad = vec![ExprId::ZERO; n_vars];
29        for i in order {
30            let Some(&g) = adj.get(&i) else { continue };
31            let id = ExprId(i as u32);
32            match self.node(id) {
33                Node::Var(n) => {
34                    if (n as usize) < n_vars {
35                        grad[n as usize] = g;
36                    }
37                }
38                Node::Lit(_) => {}
39                Node::Add(a, b) => {
40                    self.accumulate(&mut adj, a, g);
41                    self.accumulate(&mut adj, b, g);
42                }
43                Node::Mul(a, b) => {
44                    let ga = self.mul(g, b);
45                    let gb = self.mul(g, a);
46                    self.accumulate(&mut adj, a, ga);
47                    self.accumulate(&mut adj, b, gb);
48                }
49                Node::Neg(a) => {
50                    let ga = self.neg(g);
51                    self.accumulate(&mut adj, a, ga);
52                }
53                Node::Recip(a) => {
54                    // d(1/a) = -(1/a)² da
55                    let rr = self.mul(id, id);
56                    let t = self.mul(g, rr);
57                    let ga = self.neg(t);
58                    self.accumulate(&mut adj, a, ga);
59                }
60                Node::Sqrt(a) => {
61                    // d√a = da / (2√a)
62                    let two = self.mul(ExprId::TWO, id);
63                    let r = self.recip(two);
64                    let ga = self.mul(g, r);
65                    self.accumulate(&mut adj, a, ga);
66                }
67                Node::Sin(a) => {
68                    let half_pi = self.lit(std::f64::consts::FRAC_PI_2);
69                    let shifted = self.add(a, half_pi);
70                    let cos = self.sin(shifted);
71                    let ga = self.mul(g, cos);
72                    self.accumulate(&mut adj, a, ga);
73                }
74                Node::Atan2(y, x) => {
75                    // ∂/∂y = x / (x² + y²), ∂/∂x = -y / (x² + y²)
76                    let xx = self.mul(x, x);
77                    let yy = self.mul(y, y);
78                    let den = self.add(xx, yy);
79                    let r = self.recip(den);
80                    let gr = self.mul(g, r);
81                    let gy = self.mul(gr, x);
82                    let t = self.mul(gr, y);
83                    let gx = self.neg(t);
84                    self.accumulate(&mut adj, y, gy);
85                    self.accumulate(&mut adj, x, gx);
86                }
87                Node::Exp2(a) => {
88                    let ln2 = self.lit(std::f64::consts::LN_2);
89                    let t = self.mul(ln2, id);
90                    let ga = self.mul(g, t);
91                    self.accumulate(&mut adj, a, ga);
92                }
93                Node::Log2(a) => {
94                    let ln2 = self.lit(std::f64::consts::LN_2);
95                    let t = self.mul(ln2, a);
96                    let r = self.recip(t);
97                    let ga = self.mul(g, r);
98                    self.accumulate(&mut adj, a, ga);
99                }
100                Node::Select(c, a, b) => {
101                    let ga = self.select(c, g, ExprId::ZERO);
102                    let gb = self.select(c, ExprId::ZERO, g);
103                    self.accumulate(&mut adj, a, ga);
104                    self.accumulate(&mut adj, b, gb);
105                }
106            }
107        }
108        grad
109    }
110
111    fn accumulate(&mut self, adj: &mut HashMap<usize, ExprId>, at: ExprId, g: ExprId) {
112        let i = at.0 as usize;
113        let sum = match adj.get(&i) {
114            Some(&prev) => self.add(prev, g),
115            None => g,
116        };
117        adj.insert(i, sum);
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use crate::graph::ExprGraph;
124
125    /// Every op, against forward `diff`, at a few points.
126    #[test]
127    fn vjp_matches_forward_diff() {
128        let mut g = ExprGraph::new();
129        let (x, y, z) = (g.var(0), g.var(1), g.var(2));
130        let xy = g.mul(x, y);
131        let s = g.sin(xy);
132        let r = g.recip(z);
133        let sq = g.sqrt(z);
134        let at = g.atan2(y, x);
135        let e = g.exp2(x);
136        let l = g.log2(z);
137        let n = g.neg(at);
138        let sel = g.select(x, s, r);
139        let t1 = g.add(sel, sq);
140        let t2 = g.mul(t1, n);
141        let t3 = g.add(e, l);
142        let o0 = g.add(t2, t3);
143        let o1 = g.mul(xy, xy);
144        let outs = [o0, o1];
145        let seeds = [g.var(3), g.var(4)];
146        let grad = g.vjp(&outs, &seeds, 3);
147        let fwd: Vec<_> = (0..3u16).map(|k| [g.diff(o0, k), g.diff(o1, k)]).collect();
148        for pt in [[0.7, -0.4, 2.3, 1.5, -0.25], [-1.2, 0.9, 0.6, -0.5, 2.0]] {
149            for k in 0..3 {
150                let want: f64 =
151                    pt[3] * g.eval::<f64>(fwd[k][0], &pt) + pt[4] * g.eval::<f64>(fwd[k][1], &pt);
152                let got: f64 = g.eval(grad[k], &pt);
153                assert!(
154                    (got - want).abs() < 1e-12 * want.abs().max(1.0),
155                    "var {k}: {got} vs {want}"
156                );
157            }
158        }
159    }
160}
tang_expr/reverse.rs

tang_expr/
reverse.rs
