Chapter 03 — Down the Rabbit Hole: What Writing a Toy Autograd Engine Taught Me

This one started as a distraction. I was deep in a fine-tuning bug — gradients looked fine on paper but the loss wasn't moving — and instead of staring at the same training loop for a fourth hour, I opened a blank file and asked myself a question I'd been quietly avoiding for two years: what does .backward() actually do? Not the calculus — I know the calculus. What does the code do. Four hours later I had a ~150-line autograd engine and, more importantly, I understood exactly where my real bug was.

The core idea is smaller than it feels

Every autograd system is really just two things: a way to record what operations produced a value, and a way to walk that record backwards applying the chain rule. A scalar-valued node needs to remember its inputs and know how to push a gradient back to each of them:

class Value:
    def __init__(self, data, children=(), op=""):
        self.data = data
        self.grad = 0.0
        self._children = children
        self._op = op
        self._backward = lambda: None  # filled in by the op that created it
 
    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
            other.grad += out.grad
        out._backward = _backward
        return out
 
    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
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

That's it. That's the trick. Every op just needs to know its local derivative and how to distribute an incoming gradient to its inputs. The "graph" is never built explicitly as a separate data structure — it's implicit in the _children pointers, which fall naturally out of just doing normal Python arithmetic.

Backward pass is a topological sort, not magic

The part that had been a black box to me was how .backward() knows what order to visit nodes in. The answer turns out to be almost embarrassingly simple: a depth-first postorder traversal, so that by the time you process a node, everything downstream of it has already had its gradient finalized.

def backward(self):
    topo, visited = [], set()
 
    def build(node):
        if node not in visited:
            visited.add(node)
            for child in node._children:
                build(child)
            topo.append(node)
    build(self)
 
    self.grad = 1.0
    for node in reversed(topo):
        node._backward()

Seeing this made something click that three separate deep learning courses hadn't quite landed: the reason PyTorch can't compute a gradient until the full forward graph exists is that this traversal needs the graph to exist first. It's not lazily inferring order as it goes — it builds the full dependency order, then walks it in reverse. Obvious in hindsight. Not obvious from just calling .backward() for two years.

Where this actually paid off

Back to the original bug: my real model's gradients were fine at every layer except one custom masking operation, where I'd written the forward pass correctly but the backward path implicitly assumed the mask was constant — which meant a .detach() was silently needed and wasn't there. I would not have gone looking for that if I hadn't just spent an afternoon rebuilding, by hand, the exact mechanism that was misbehaving. Toy implementations are usually framed as pedagogy — "build it to learn it" — but the underrated version of that advice is narrower and more useful: build the toy version of the specific mechanism your real bug lives in, not the whole system. I didn't need a toy transformer. I needed a toy .backward(), and that was exactly the rabbit hole worth going down.