xtrain/crates/xtrain-autodiff/src/tape.rs

//! Tape-based reverse-mode autograd (Phase T4).
//!
//! A [`Var`] is a reference-counted node wrapping a value [`Tensor`], an optional
//! accumulated gradient, and — for non-leaf nodes — the parents it was computed
//! from plus a backward closure. Forward ops (see [`crate::ops`]) build these
//! nodes; [`Var::backward`] walks the graph in reverse topological order and
//! pushes gradients to parents, **summing** on fan-out (a tensor consumed by
//! several ops).
//!
//! The graph is dynamic (define-by-run) and the design favours clarity over
//! speed: each op synchronizes its own kernels (T3 has no streams yet), and
//! gradient accumulation is an explicit elementwise add.

#![cfg(not(no_cuda))]

use std::cell::RefCell;
use std::rc::Rc;
use xtrain_tensor::Tensor;

/// Backward closure: given this node's accumulated grad and its parents, compute
/// and accumulate each parent's gradient contribution.
pub type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;

pub struct VarNode {
    pub value: Tensor,
    pub grad: Option<Tensor>,
    parents: Vec<Var>,
    backward: Option<BackwardFn>,
}

/// A node in the autograd tape. Cheap to clone (bumps the `Rc`); clones share
/// the same underlying node, which is how fan-out is detected.
#[derive(Clone)]
pub struct Var(Rc<RefCell<VarNode>>);

impl Var {
    /// A leaf node (parameter / input). After `backward`, its `.grad()` holds the
    /// accumulated gradient of the loss w.r.t. this tensor.
    pub fn leaf(value: Tensor) -> Self {
        Var(Rc::new(RefCell::new(VarNode {
            value,
            grad: None,
            parents: Vec::new(),
            backward: None,
        })))
    }

    /// Build a non-leaf node from a forward `value`, its `parents`, and a
    /// `backward` closure. Used by the op constructors in [`crate::ops`] and by
    /// callers that compose custom nodes (e.g. a loss reduction or a transpose
    /// the built-in op set doesn't cover).
    pub fn from_op(value: Tensor, parents: Vec<Var>, backward: BackwardFn) -> Self {
        Var(Rc::new(RefCell::new(VarNode {
            value,
            grad: None,
            parents,
            backward: Some(backward),
        })))
    }

    /// Clone of the value tensor (cheap: storage is `Arc`-shared).
    pub fn value(&self) -> Tensor {
        self.0.borrow().value.clone()
    }

    /// The accumulated gradient, if any (populated after `backward`).
    pub fn grad(&self) -> Option<Tensor> {
        self.0.borrow().grad.clone()
    }

    /// Clear the accumulated gradient. Call on every parameter between training
    /// steps so the next `backward` accumulates from zero (grads SUM otherwise).
    pub fn zero_grad(&self) {
        self.0.borrow_mut().grad = None;
    }

    /// Overwrite this node's value tensor in place. Used by the optimizer to
    /// apply a parameter update (`p ← p − lr·grad`) while keeping the leaf's
    /// identity stable across steps.
    pub fn set_value(&self, value: Tensor) {
        self.0.borrow_mut().value = value;
    }

    /// Pointer identity, used to dedup nodes during the topological sort.
    fn id(&self) -> *const RefCell<VarNode> {
        Rc::as_ptr(&self.0)
    }

    /// Accumulate `g` into this node's grad slot (SUM — the fan-out rule).
    fn accumulate(&self, g: Tensor) {
        let mut node = self.0.borrow_mut();
        match node.grad.take() {
            None => node.grad = Some(g),
            Some(existing) => node.grad = Some(existing.add(&g)),
        }
    }

    /// Reverse-mode backward from this node (treated as a scalar loss: its
    /// upstream grad is seeded to ones). Populates `.grad` on every node that
    /// transitively feeds it.
    ///
    /// `loss` must be a scalar (single element) so the seed `dL/dL = 1` is
    /// unambiguous, matching the `L = sum(W∘out)` grad-check convention.
    pub fn backward(&self) {
        assert_eq!(
            self.value().numel(),
            1,
            "backward() expects a scalar loss; got shape {:?}",
            self.value().shape()
        );

        // 1. Topological order (post-order DFS), parents before children.
        let mut topo: Vec<Var> = Vec::new();
        let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
        build_topo(self, &mut topo, &mut visited);

        // 2. Seed the loss gradient with ones.
        let seed = ones_like(&self.value());
        self.accumulate(seed);

        // 3. Walk in reverse: each node hands its grad to its parents' closures.
        for var in topo.iter().rev() {
            let (grad, parents, backward) = {
                let node = var.0.borrow();
                (
                    node.grad.clone(),
                    node.parents.clone(),
                    node.backward.is_some(),
                )
            };
            let grad = match grad {
                Some(g) => g,
                None => continue, // node didn't contribute to the loss
            };
            if backward {
                // Borrow the closure out, run it, then drop the borrow.
                let node = var.0.borrow();
                let bw = node.backward.as_ref().unwrap();
                bw(&grad, &parents);
            }
        }
    }

    /// Drop the parents/closure so the graph can be freed, keeping value+grad.
    /// (Not needed for tests; provided for completeness of the engine.)
    pub fn detach_graph(&self) {
        let mut node = self.0.borrow_mut();
        node.parents.clear();
        node.backward = None;
    }

    /// Distribute `g` to a parent (called from op backward closures). Every node
    /// accumulates its grad: intermediates need it for the chain rule, leaves
    /// expose it as the result. This SUM is what makes fan-out correct.
    pub fn push_grad(parent: &Var, g: Tensor) {
        parent.accumulate(g);
    }
}

/// Post-order DFS: append a node only after all its parents, dedup by identity.
fn build_topo(var: &Var, topo: &mut Vec<Var>, visited: &mut Vec<*const RefCell<VarNode>>) {
    if visited.contains(&var.id()) {
        return;
    }
    visited.push(var.id());
    let parents = var.0.borrow().parents.clone();
    for p in &parents {
        build_topo(p, topo, visited);
    }
    topo.push(var.clone());
}

/// A ones tensor matching `t`'s shape/device (the backward seed for a scalar).
fn ones_like(t: &Tensor) -> Tensor {
    let host = vec![1.0f32; t.numel()];
    Tensor::from_slice(&host, t.shape()).to_device(t.device())
}