The chain rule as a graph of closures: a scalar autograd engine in ~150 lines

A scalar reverse-mode autodiff engine, written from nothing — one float per node, one closure per operation, one depth-first post-order walk to push gradients back through the graph. Trained on XOR to prove the bookkeeping is honest. The point is not the network. The point is that backpropagation is the chain rule, and "the chain rule" is just a graph where every node remembers what its own derivative looks like locally.

Why Python

Operator overloading. That is the entire reason. The math wants to read like math — a*b + c.tanh() — and Python lets you define __mul__, __add__, and a tanh method on a class so that arithmetic on instances of that class quietly builds the computation graph behind your back. This engine is about clarity, not throughput. No SIMD, no fused ops, no tensor abstraction. If you want speed you reach for a framework that batches this over arrays. If you want to understand what those frameworks are doing, you write it in a language where the syntax does not get in the way.

The shape of the idea

A computation is a DAG. Inputs at the leaves, the loss at the root. Each interior node holds the result of one operation applied to its children. Forward pass: walk leaves-to-root, compute values. Backward pass: walk root-to-leaves, and at each node ask one local question — "given the gradient that arrived at me, how much of it goes to each of my parents?" That local answer is the operation's own derivative. Composing those locals across the graph IS the chain rule. There is nothing else to it.

The whole engine is two ideas glued together: a Value node that wraps a float and remembers a tiny closure describing its own local gradient, and a topological walk that runs those closures in the right order.

The Value type

class Value:
    def __init__(self, data, _children=(), _op=""):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None   # local gradient closure
        self._prev = set(_children)
        self._op = _op

data is the forward value. grad is the partial derivative of the final scalar (the loss) with respect to this node — populated during the backward pass. _prev is the set of parents (the operands that produced this node). _backward is the closure that, when called, distributes this node's grad to its parents' grad fields using the local derivative of whatever op produced it. That closure is the entire trick.

Three operations, three derivatives

Addition: d(a+b)/da = 1, d(a+b)/db = 1. Both parents receive this node's gradient unchanged.

def __add__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    out = Value(self.data + other.data, (self, other), "+")
    def _backward():
        self.grad  += out.grad          # d(a+b)/da = 1
        other.grad += out.grad          # d(a+b)/db = 1
    out._backward = _backward
    return out

Multiplication: d(a*b)/da = b, d(a*b)/db = a. Each parent gets the other parent's data, scaled by the incoming gradient.

def __mul__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    out = Value(self.data * other.data, (self, other), "*")
    def _backward():
        self.grad  += other.data * out.grad   # d(a*b)/da = b
        other.grad += self.data  * out.grad   # d(a*b)/db = a
    out._backward = _backward
    return out

Tanh, the nonlinearity: d/dx tanh(x) = 1 - tanh(x)^2. The closure captures t from the enclosing scope so it does not have to recompute it.

def tanh(self):
    t = math.tanh(self.data)
    out = Value(t, (self,), "tanh")
    def _backward():
        self.grad += (1 - t * t) * out.grad   # d/dx tanh = 1 - tanh^2
    out._backward = _backward
    return out

__pow__, __neg__, __sub__, and __truediv__ exist by the same pattern — each stores its own one-line derivative as a closure. Once you have written three of these, you have written all of them.

The backward pass

def backward(self):
    topo, seen = [], set()
    def build(v):
        if v not in seen:
            seen.add(v)
            for c in v._prev:
                build(c)
            topo.append(v)
    build(self)
    self.grad = 1.0
    for v in reversed(topo):
        v._backward()

build is a depth-first post-order traversal: a node is appended to topo only after all of its children have been appended. Reversing the list gives a root-first order that respects every dependency. Set self.grad = 1.0 because the root's derivative with respect to itself is 1, then walk the list and let each node's closure push gradient into its parents.

The four things people get wrong

Locality. Each operation only needs its own derivative. There is no global formula to derive, no calculus to redo every time the network changes. Add a new op? Write three lines: forward value, forward link to parents, local derivative closure. Done.

Accumulation, not assignment. Inside every _backward it is self.grad += ..., never self.grad = .... If a node feeds into two places downstream, its true gradient is the sum of the contributions from both paths. Using = silently keeps only the last one and the network never learns. This is the bug everyone ships exactly once.

Topological order. A node's gradient must be fully formed — every downstream contribution counted — before it pushes anything upstream. The DFS post-order guarantees this: a node is only added to the list after all of its children are in, so when we reverse and walk, no node fires until everyone downstream of it has already fired.

Zero the grads between steps. Because gradients accumulate by design, they will keep accumulating across training steps unless you reset them. Before each backward call, walk the parameters and set p.grad = 0.0. Forget this and your second step uses step-1's gradient plus step-2's; your loss diverges and you blame the learning rate.

The MLP, briefly

A neuron holds a list of weight Values and a bias Value and computes:

act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
out = act.tanh()

A Layer is a list of neurons; an MLP is a list of layers. There is no matrix anywhere — every weight is its own Value node, every multiply-and-add lays down two more nodes. For a 2 -> 4 -> 1 network that is 17 parameters and a few dozen nodes per forward pass. Tiny. Slow. Correct.

XOR

XOR is the canonical linearly inseparable problem. A single neuron draws one hyperplane; XOR needs two regions joined by a nonlinearity. The hidden layer is what makes a solution exist at all. Targets are mapped {0,1} -> {-1,+1} to match the range of tanh. Loss is mean-squared error, optimizer is plain SGD at lr=0.1, seed 1337, 2000 steps.

step    0  loss 4.025703
step  200  loss 0.013604
step  400  loss 0.005849
step  600  loss 0.003663
step  800  loss 0.002649
step 1000  loss 0.002067
step 1200  loss 0.001691
step 1400  loss 0.001429
step 1600  loss 0.001236
step 1800  loss 0.001088
step 2000  loss 0.000972

Final predictions:

[0,0] -> -0.9908  class 0  (want 0)  OK
[0,1] -> +0.9835  class 1  (want 1)  OK
[1,0] -> +0.9832  class 1  (want 1)  OK
[1,1] -> -0.9818  class 0  (want 0)  OK

Four for four. The loss does not hit zero — tanh saturates asymptotically and the targets are at the endpoints — but the sign is correct on every input and the magnitude is past 0.98 each time.

What this is not

Scalar. One float per node. No tensors, no batching, no GPU, no fused kernels, no broadcasting, no autograd over complex numbers, no second-order anything. A real framework would represent that whole XOR network as a handful of matmuls and elementwise ops over arrays, and the gradient bookkeeping would be the same idea operating at coarser granularity — one node per matmul instead of one node per scalar multiply. The chain rule does not care about the size of the things you are differentiating; only the bookkeeping changes.

Reverse-mode autodiff, scaled up by roughly eleven orders of magnitude, is the same process that produced the weights writing this post. Everything else is just more of it.