oxmera_tensor/
autograd.rs

1//! Reverse-mode autograd plumbing: the tape nodes tensors carry, gradient
2//! accumulation, and the recording switch.
3//!
4//! The differentiable *rules* (VJPs) live next to the ops that create them
5//! (`crate::ops`); this module owns the graph mechanics only.
6
7use std::cell::Cell;
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex};
10
11use oxmera_core::Result;
12
13use crate::tensor::Tensor;
14
15/// One recorded operation: the inputs it consumed and the vector-Jacobian
16/// product mapping the output gradient to input gradients (aligned with
17/// `inputs`; `None` for non-differentiable inputs such as indices).
18pub struct GradFn {
19    /// The tensors the op consumed, in order.
20    pub inputs: Vec<Tensor>,
21    /// The VJP: output gradient in, one optional gradient per input out.
22    #[allow(clippy::type_complexity)]
23    pub vjp: Box<dyn Fn(&Tensor) -> Result<Vec<Option<Tensor>>> + Send + Sync>,
24}
25
26impl std::fmt::Debug for GradFn {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("GradFn")
29            .field("inputs", &self.inputs.len())
30            .finish()
31    }
32}
33
34/// The autograd state a tracked tensor carries.
35#[derive(Debug)]
36pub struct AutogradMeta {
37    /// Whether gradients accumulate on this tensor during `backward`.
38    pub(crate) requires_grad: bool,
39    /// The accumulated gradient, if any backward pass has reached it.
40    pub(crate) grad: Mutex<Option<Tensor>>,
41    /// How this tensor was computed; `None` for leaves.
42    pub(crate) grad_fn: Option<GradFn>,
43}
44
45thread_local! {
46    static RECORDING: Cell<bool> = const { Cell::new(true) };
47}
48
49/// Whether ops on this thread currently record onto the tape.
50pub fn is_recording() -> bool {
51    RECORDING.with(Cell::get)
52}
53
54/// Run `f` with tape recording disabled — the inference/`no_grad` context.
55///
56/// Nested calls are fine; recording resumes when the outermost guard ends.
57pub fn no_grad<R>(f: impl FnOnce() -> R) -> R {
58    let _guard = NoGradGuard::new();
59    f()
60}
61
62/// RAII guard that disables tape recording until dropped.
63pub struct NoGradGuard {
64    previous: bool,
65}
66
67impl NoGradGuard {
68    /// Disable recording on this thread until the guard drops.
69    pub fn new() -> Self {
70        let previous = RECORDING.with(|r| r.replace(false));
71        Self { previous }
72    }
73}
74
75impl Default for NoGradGuard {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl Drop for NoGradGuard {
82    fn drop(&mut self) {
83        let previous = self.previous;
84        RECORDING.with(|r| r.set(previous));
85    }
86}
87
88fn meta_id(meta: &Arc<AutogradMeta>) -> usize {
89    Arc::as_ptr(meta) as usize
90}
91
92/// Reverse-topological gradient propagation from `root`, seeding with
93/// `seed`. Called by [`Tensor::backward`].
94pub(crate) fn run_backward(root: &Tensor, seed: Tensor) -> Result<()> {
95    // Everything backward computes is bookkeeping, not new graph.
96    no_grad(|| run_backward_inner(root, seed))
97}
98
99fn run_backward_inner(root: &Tensor, seed: Tensor) -> Result<()> {
100    let Some(root_meta) = root.autograd_meta() else {
101        return Ok(());
102    };
103
104    // Post-order DFS for a topological order over the tape.
105    let mut order: Vec<Tensor> = Vec::new();
106    let mut visited: HashMap<usize, ()> = HashMap::new();
107    let mut stack: Vec<(Tensor, bool)> = vec![(root.clone(), false)];
108    while let Some((t, expanded)) = stack.pop() {
109        let Some(meta) = t.autograd_meta() else {
110            continue;
111        };
112        let id = meta_id(&meta);
113        if expanded {
114            order.push(t);
115            continue;
116        }
117        if visited.contains_key(&id) {
118            continue;
119        }
120        visited.insert(id, ());
121        stack.push((t.clone(), true));
122        if let Some(gf) = &meta.grad_fn {
123            for input in &gf.inputs {
124                stack.push((input.clone(), false));
125            }
126        }
127    }
128
129    let mut pending: HashMap<usize, Tensor> = HashMap::new();
130    pending.insert(meta_id(&root_meta), seed);
131
132    for t in order.into_iter().rev() {
133        let meta = t.autograd_meta().expect("ordered tensors carry meta");
134        let id = meta_id(&meta);
135        let Some(grad) = pending.remove(&id) else {
136            continue;
137        };
138
139        if meta.requires_grad {
140            let mut slot = meta.grad.lock().expect("grad mutex poisoned");
141            *slot = Some(match slot.take() {
142                Some(existing) => existing.add(&grad)?,
143                None => grad.clone(),
144            });
145        }
146
147        if let Some(gf) = &meta.grad_fn {
148            let input_grads = (gf.vjp)(&grad)?;
149            debug_assert_eq!(input_grads.len(), gf.inputs.len());
150            for (input, ig) in gf.inputs.iter().zip(input_grads) {
151                let (Some(im), Some(ig)) = (input.autograd_meta(), ig) else {
152                    continue;
153                };
154                let iid = meta_id(&im);
155                let accumulated = match pending.remove(&iid) {
156                    Some(existing) => existing.add(&ig)?,
157                    None => ig,
158                };
159                pending.insert(iid, accumulated);
160            }
161        }
162    }
163    Ok(())
164}
oxmera_tensor/autograd.rs

oxmera_tensor/
autograd.rs
