Autograd

In short: a Variable is a tensor plus a gradient slot and a pointer to the Node that produced it. Each differentiable op computes its value with the plain tensor op and calls MakeResult, which records a Node holding a backward closure and edges to its inputs' Nodes. Leaves get an AccumulateGrad Node. Backward counts incoming edges, runs each Node once all of its gradient has arrived, sums gradients from different paths, and frees saved tensors unless asked to keep them. User-defined ops subclass Function, and GradCheck compares every gradient with central differences.

The headers link derivatives and gradients, the chain rule, backpropagation and computation graphs on /ml-models. The values inside a Variable are ordinary tensors (Tensors and dtypes), and the matrix product is the one on Linear algebra.

Variable

Variable is a handle to a shared VariableImpl (variable.cc L20–30): the value, the gradient, a requires_grad flag, the producing Node (grad_fn), a weak pointer to the leaf's gradient accumulator, and any hooks. Copying a Variable shares all of it.

Recording the graph

The graph is built while the forward pass runs ("define by run"), so ordinary C++ control flow decides its shape (variable.h L4–21). Every differentiable op, built-in or user-defined, ends in MakeResult. It wraps the computed value in a new Variable, then looks up each input's GradientEdge (L241–251): the input's grad_fn, the AccumulateGrad Node of a leaf that requires grad (created on first use and shared while it lives), or null. If grad mode is off or every edge is null, the result is a plain leaf and the backward closure is dropped. Otherwise the result gets a FunctionNode holding the closure, the edges and the input shapes.

A Node maps the gradient of its one output to one gradient per input; next_functions()[i] is the Node that receives input i's gradient. Edges point from outputs toward inputs, so the graph is owned from the top: a Variable keeps its grad_fn alive, which keeps every Node beneath it alive. A leaf holds its accumulator only weakly, so it does not keep a graph alive after the output is gone.

The backward pass

RunBackward (variable.cc L90–139) works in two phases:

1. Count. A breadth-first walk from the root counts, for every Node, how many edges lead into it. A Node may run only after that many gradient contributions have arrived, which is a topological order without sorting.
2. Run. Starting from the root with the seed, it pops a ready Node, runs its pre-hooks and backward function (Apply), and checks that it returned one f32 gradient per input with the input's recorded shape. Each gradient is added into the input's pending slot, and an input whose count reaches zero becomes ready. If no path delivered a gradient to a Node, it is skipped, but its inputs' counts still drop.

Tracing the figure: SumAll runs on 1 and passes 1 to Add. Add passes 1 to Mul and to AccumulateGrad(a), whose count drops from 2 to 1, so it waits. Mul sends g·b = 3 to a (pending 1 + 3 = 4, count 0, ready) and g·a = 2 to b. The accumulators then run once each and store 2 and 4.

AccumulateGrad (L70–88) runs the leaf's hooks, then adds into grad(). The first gradient is cloned, because it may be shared with another branch and the leaf's owner is allowed to change grad() in place. Gradients accumulate across calls until ZeroGrad, as in PyTorch.

Unless retain_graph is set, every Node visited is then released: a FunctionNode drops its closure and with it the tensors it captured (L42–68). A second Backward through the same graph aborts with a message that says to pass retain_graph = true.

Grad mode

Recording is on by default and tracked per thread (a thread_local flag, L271–277). NoGradGuard turns it off for its scope and restores the previous value, so guards nest. Under a guard, MakeResult returns plain leaves and records nothing, which is how evaluation and parameter updates avoid building a graph.

The differentiable ops

ops.h mirrors rml/core/ops.h and linalg::MatMul. Each op computes its value with the tensor op, then records a closure that captures only what its derivative needs: shapes for Add, both operands for Mul, the output for Softmax. The header comment beside each op states its derivative, with g the gradient arriving from the output. A closure skips any input whose needs_grad flag is false.

Gradients of broadcast operands

When an operand was broadcast, each of its elements fed every output element along the stretched axes, so its gradient is the sum over those axes. internal::SumTo (ops.cc L24–36) undoes a broadcast: it sums away the extra leading axes, then sums with keep_dim every axis that is 1 in the target shape but not in the gradient. Adding a bias of shape [4] to activations of shape [32, 4] therefore gives the bias a gradient of shape [4], the column sums of g. See broadcasting for the forward rule.

MatMul

MatMul's backward (ops.cc L289–310) promotes vector operands to matrices exactly as the forward pass did, reshapes g to match, and computes dA = g·Bᵀ and dB = Aᵀ·g as batched products. It sums them over broadcast batch axes with SumTo and reshapes back to each operand's original shape. Both products run on the threaded GEMM.

User-defined ops: Function

Function works like torch.autograd.Function. A subclass computes the value in Forward, keeps what Backward will need as members, and returns one gradient per input from Backward. Apply (function.cc L15–31) runs Forward under a NoGradGuard, so its internals record nothing, then records a single Node named after the Function. When the graph is released, the Node's closure (and so the Function with its saved members) is freed.

rml uses this where one fused backward is simpler or cheaper than a chain of small ops: the im2col unfold, patch rows and max pooling in nn/conv.cc, and sigmoid cross-entropy in losses/classification.cc (see Neural network layers).

Gradient checking

rml's rule is that every differentiable op gets a numerical gradient check. grad_check.h compares autograd's gradients with central differences, ∂L/∂xᵢ ≈ (L(x + ε eᵢ) − L(x − ε eᵢ)) / 2ε, whose error shrinks as ε².

• The function may return any shape. It is reduced to L = Σ w ⊙ f(inputs) with fixed pseudo-random weights w in [−1, 1), so every output element counts, and swapping two output elements would change the gradient (grad_check.cc L22–41). The weights come from raw mt19937 bits because the standard distributions differ between standard libraries.
• GradCheckVariables (L45–82) runs one Backward for the analytic gradient, then, under NoGradGuard, nudges each element by ±ε in place, re-evaluates, and restores it. An element passes when |analytic − numeric| ≤ atol + rtol·|numeric|, with defaults ε = 10⁻³, atol = 10⁻³ and rtol = 10⁻². The first failure comes back as a Status naming the input and element.
• GradCheck does the same for fresh copies of some tensors. GradCheckVariables is for Variables that already exist, such as a module's parameters, which fn reads by capture.

Tests

How these run under sanitizers is on Testing.