The Autograd Engine: Rules Belong to the Nodes, Order to the Topological Sort

The first three parts were about understanding backpropagation. This one is about turning it into code that runs by itself.

There is no new mathematics here. The chain rule, δ and multi-path summation have all been derived already; the difficulty is entirely engineering.

Where numerical differentiation ended up

Rather than starting from symbolic differentiation, start from the crudest method — perturbing the input:

f'(x) ≈ [ f(x+h) − f(x) ] / h

For f(x) = 3x² − 4x + 5 at x = 3, taking h = 0.001 gives 14.003; the analytic answer 6x − 4 = 14. It checks out.

But take h from 0.1 down to 1e-14 and the result gets more accurate, and then suddenly falls apart. Once h is small enough, f(x+h) and f(x) are nearly equal in floating point, and subtracting them eats the significant digits.

So numerical differentiation cannot be used for training, for two reasons:

• It is slow: every parameter needs its own forward pass. 13,002 parameters means 13,002 passes
• It is unstable: past a certain point, a smaller h makes the accuracy worse

It has not disappeared, though — it just has a different job now: gradient checking. Once backpropagation is written, numerical differentiation verifies that its gradients are right.

Value: making every number remember where it came from

The core idea: after d = a*b + c runs, d should be carrying the information “I came from adding the result of a*b to c”. Wire all the numbers together and you have a computation graph.

Four fields:

The graph is built by hand — Python does nothing to help. What makes it work is that _prev holds object references, not copies of numbers; follow the references down and you have a tree. Anything with an empty _prev is a leaf (an input or a parameter).

Operator overloading is the other half: because __add__ and __mul__ are implemented, you can write d = a*b + c, an expression identical to ordinary maths, that quietly builds a graph behind your back. PyTorch’s tensors work by overloading the same methods.

The animation below takes one tanh neuron apart into 16 steps: building the graph → topological sort → ten calls to _backward.

Filling it in by hand: two rules are enough

Before automating anything, do it by hand once. Only two rules are needed:

• An add node passes the gradient to both children unchanged (local derivative 1)
• A multiply node multiplies the gradient by the other factor before passing it on

All ten nodes filled in:

Two things worth noting:

w2.grad = 0, because x2 = 0. Part 2’s rule — if any one of the three factors is near 0, that weight barely needs to move — this is its numeric form: the neuron below carries no signal, so whatever you do to the weight attached to it gets multiplied by 0.

The four add nodes all have grad 0.5, undiminished the whole way. Addition’s local derivative is 1, and multiplying by 1 costs nothing. It is also why a bias’s gradient always equals its layer’s δ.

There is a trap in doing it by hand: to refer to the intermediate nodes, n = x1*w1 + x2*w2 + b has to be split into five statements with a name for every intermediate result. Create s1 = w1 * x1 partway through and everything goes wrong — that is a brand new object, not the one on the graph. The manual approach is simply unusable on a real network.

Separating local rules from global scheduling

The key observation: those two rules have nothing to do with the graph as a whole, only with what operation I am. A multiply node always passes gradients by the multiply rule, wherever in the graph it sits.

So push the rule down into the node itself:

This is the skeleton of every automatic differentiation framework.

Closures: storing the rule in the node itself

def __mul__(self, 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     # store the function, do not call it
    return out

Three things matter:

• The order cannot be reversed: create out first, then define _backward, because the body refers to out.grad
• Nothing is evaluated at definition time: out.grad is still 0 right now, and is only read when the backward pass arrives. A closure captures the variable itself, not a snapshot
• Every node binds its own three objects: ten thousand multiply nodes means ten thousand _backward functions, each remembering its own self / other / out, none interfering with the others

The closure binds its arguments in advance, so calling it takes no arguments — which is what lets the loop inside backward() stay as clean as it is.

The topological sort: where append sits decides correctness

def backward(self):
    topo, visited = [], set()
    def build(v):
        if v not in visited:
            visited.add(v)
            for child in v._prev:
                build(child)
            topo.append(v)      # children first, then self
    build(self)
    self.grad = 1.0             # open the tap
    for node in reversed(topo):
        node._backward()

topo.append(v) sits after the recursion, and that single fact decides whether the algorithm is correct. Finish every child before appending yourself and children are guaranteed to precede their parent in the list; walking it in reverse then goes from output back to inputs.

self.grad = 1.0 is the only number in the whole chain nobody has to compute (∂o/∂o = 1); everything else is passed along internally.

A leaf’s _backward is lambda: None, so the loop can call every node indiscriminately, without ever checking whether something is a leaf.

Why it has to be +=

Getting it wrong on purpose once makes it obvious:

a = Value(3.0)
b = a + a
b.backward()
# with =   → a.grad = 1.0  ❌
# with +=  → a.grad = 2.0  ✅

The reason is that in a + a, self and other are the same object, so both lines of _backward write to the same a.grad. With =, the second write overwrites the first.

And since b = a + a = 2a, ∂b/∂a = 2 is the correct answer.

This plus sign is the code form of the Σ from part 3. There it was “h1 has two downstreams, so there are two terms to add”; here a is used twice, so there are two +=. The elegant part is that _backward only handles its own contribution — the accumulation happens on its own through += — and whoever writes the code never has to work out whether a node is used in more than one place.

The matching cost: grad is never cleared automatically. A second backward() accumulates on top of the first. A training loop has to zero them every round — which is exactly why optimizer.zero_grad() exists in PyTorch.

Filling out the operators: primitive and composite

The new primitives (they create a node, so they need a _backward):

The composites (they only forward, and need no _backward):

def __truediv__(self, other): return self * other**-1
def __neg__(self):            return self * -1
def __sub__(self, other):     return self + (-other)
def __rmul__(self, other):    return self * other
def __radd__(self, other):    return self + other

The test is simple: if Value(...) appears in the code, it needs a _backward; if it only does return self * ..., it does not. There is no “subtract node” anywhere in the graph, just a multiply and an add; division picks up the correct gradient for free through __pow__(-1).

Three fewer _backward written is three fewer places to get one wrong.

When __rmul__ gets called follows from how Python evaluates x * y: it tries type(x).__mul__(x, y) first and stops if that works; if it returns NotImplemented, it tries type(y).__rmul__(y, x). The r stands for reflected — the left side cannot handle it, so the right side is asked whether it wants to. Which is what makes 2 * n work.

__add__ and __mul__ also have to accept plain numbers, so they open with:

other = other if isinstance(other, Value) else Value(other)

Expanding tanh: you choose the node granularity

tanh was used as a single operator above, with its derivative 1 − t² worked out by hand. But it can also be assembled from primitives:

e   = (2*n).exp()
num = e - 1
den = e + 1
o   = num / den
o.backward()

The leaf gradients are identical to the packed version (−1.5 / 1.0 / 0.5 / 0), but the graph goes from 5 nodes to a dozen or so.

The part worth watching in the animation is the fork at e: it is used by both num and den, so on the way back two _backward calls write to e.grad in turn — −0.1036 plus +0.1464 accumulating into +0.0429. The += from the previous section, in the flesh.

Which gives the conclusion: automatic differentiation does not dictate the granularity at which you define operators.

This is why fused operators like F.softmax and F.layer_norm exist in PyTorch — they could have been assembled from primitives, and packing them is for speed and numerical stability.

Automatic differentiation guarantees mathematical correctness, not numerical stability. Write sigmoid out as 1 / (1 + (-x).exp()) and at x = -1000, exp(1000) overflows outright. The gradient is right; you just cannot get a number out of it.

In passing: eˣ cannot be an activation function

• Unbounded output: x=10 → 22026, x=50 → 5×10²¹; stack a few layers and it overflows to inf/NaN
• Its derivative is also eˣ: the gradient explodes in step — the opposite extreme to the vanishing gradients of part 2

It is all over neural networks, but always held in check by a division: sigmoid’s denominator, tanh’s numerator and denominator, softmax’s division by the total. exp is implemented here to take tanh apart, not to be used as an activation.

The complete code

The sections above take it apart piece by piece; here is a version that runs as it stands.

Value

import math

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = data
        self._prev = set(_children)
        self._op = _op
        self.grad = 0.0
        self.label = label
        self._backward = lambda: None

    def __repr__(self):
        return f"Value(data={self.data})"

    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 += 1.0 * out.grad
            other.grad += 1.0 * 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

    def tanh(self):
        t = (math.exp(2*self.data) - 1) / (math.exp(2*self.data) + 1)
        out = Value(t, (self,), 'tanh')

        def _backward():
            self.grad += (1 - t**2) * out.grad
        out._backward = _backward

        return out

    def exp(self):
        out = Value(math.exp(self.data), (self,), 'exp')

        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward

        return out

    def __pow__(self, other):
        assert isinstance(other, (int, float))
        out = Value(self.data ** other, (self,), f'**{other}')

        def _backward():
            self.grad += other * (self.data ** (other-1)) * out.grad
        out._backward = _backward

        return out

    # Composites below: they only forward, create no node, and need no _backward
    def __truediv__(self, other): return self * other**-1
    def __neg__(self):            return self * -1
    def __sub__(self, other):     return self + (-other)
    def __rmul__(self, other):    return self * other
    def __radd__(self, other):    return self + other

    def backward(self):
        topo = []
        visited = set()

        def build(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build(child)
                topo.append(v)

        build(self)

        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

        print([v.label for v in topo])

That print at the end is for debugging — it dumps the topological order, which is what tells the two examples below apart.

Drawing the graph

draw_dot needs graphviz. It is only there to look at the graph and has nothing to do with differentiation:

from graphviz import Digraph

def trace(root):
    """Walk the whole graph along _prev, collecting every node and edge"""
    nodes, edges = set(), set()
    def walk(v):
        if v not in nodes:
            nodes.add(v)
            for child in v._prev:
                edges.add((child, v))
                walk(child)
    walk(root)
    return nodes, edges

def draw_dot(root):
    dot = Digraph(format='svg', graph_attr={'rankdir': 'LR'})
    nodes, edges = trace(root)

    for n in nodes:
        uid = str(id(n))
        label = getattr(n, 'label', '')
        dot.node(name=uid,
                 label=f"{{ {label} | data {n.data:.4f} | grad {n.grad:.4f} }}",
                 shape='record')
        if n._op:
            dot.node(name=uid + n._op, label=n._op)
            dot.edge(uid + n._op, uid)

    for a, b in edges:
        dot.edge(str(id(a)), str(id(b)) + b._op)

    return dot

Two examples

The front half is the same in both; only that final tanh step differs.

x1 = Value(2.0, label='x1'); x2 = Value(0.0, label='x2')
w1 = Value(-3.0, label='w1'); w2 = Value(1.0, label='w2')
b  = Value(6.8813735870195432, label='b')

x1w1 = x1*w1;        x1w1.label = 'x1*w1'
x2w2 = x2*w2;        x2w2.label = 'x2*w2'
s    = x1w1 + x2w2;  s.label = 'x1w1+x2w2'
n    = s + b;        n.label = 'n'

o = n.tanh();        o.label = 'o'      # packed
o.backward()
draw_dot(o)

The topological order printed is 10 nodes (from one particular run):

['b', 'x1', 'w1', 'x1*w1', 'x2', 'w2', 'x2*w2', 'x1w1+x2w2', 'n', 'o']

Swap the last part for the expanded version:

e   = (n * 2).exp();  e.label = 'e'      # 2*n works too, via __rmul__
num = e - 1;          num.label = 'num'
den = e + 1;          den.label = 'den'
o   = num / den;      o.label = 'o'
o.backward()

The order becomes 18 nodes:

['', '', 'b', 'w1', 'x1', 'x1*w1', 'w2', 'x2', 'x2*w2',
 'x1w1+x2w2', 'n', '', 'e', 'den', '', '', 'num', 'o']

Five of the eight extra entries have an empty label — those are the intermediates nobody named: the constants 2, 1 and -1, each wrapped into a leaf Value, plus the n*2 multiply node and the den**-1 power node.

You will most likely not reproduce the order of either list. _prev is a set, so which sibling comes first depends on object ids and can differ between runs and between machines. The only thing a topological sort guarantees is that children precede their parent — and that is all it needs to guarantee, because walking it in reverse only requires that everything downstream of a node is already done. The node counts (10 and 18) are stable.

The value of b, 6.8813735870195432, was not picked at random: it makes n = 0.8813735870 and o = tanh(n) = 0.7071067812 — that is √2/2 — so n.grad = 1 − o² = 0.5 comes out a round number, which makes checking against hand calculation easy.

Run both versions and the leaf gradients are identical.

How this maps onto the first three parts

The whole part in one sentence: store “how each operation passes gradients” in the node itself, and hand “what order to pass them in” to a topological sort. Separate the two and you have an engine that handles any expression — as long as the primitives’ derivatives are right, any composite function assembled from them gets its gradients for free.

The engine is written, but it still only differentiates expressions; it is not a neural network yet.

Building the Neuron / Layer / MLP classes on top of it is the next part.