Build Your Own Autograd — A Mini PyTorch in 100 Lines

Start from a single Value class, add operator overloading, topological ordering, and gradient accumulation, then put a neural network on top and train it. Once you have seen the reasons behind each design choice, zero_grad() and retain_graph stop being trivia to memorize.

What reading gives you, and what building gives you

loss.backward() is nothing more than the chain rule applied mechanically — that argument is laid out in How Automatic Differentiation Works. But there is a gap between nodding along to an explanation and watching a class you wrote drop the right number into w.grad.

Here we build a working automatic differentiation engine that handles a single scalar at a time, stopping at each design decision to ask why it has to be that way. The finished thing is a little over 100 lines. Then we stack a small neural network on it and train until the loss comes down. Karpathy's micrograd is the best-known implementation in this style, and the skeleton below is from that lineage.

What building it actually answers are the questions that nagged you as a user. Why does zero_grad() have to be called every step? Why do gradients accumulate instead of being assigned? Why does calling backward() twice on the same graph raise an error? From the inside, none of these are arbitrary rules. They are the only way the thing could have been built.

A node is a small delivery worker

Every value that appears in the computation becomes a Value object. 2.0 is a Value. The result of w * x is a Value. So is the final loss.

Each Value carries three things: its own number, the gradient that has arrived at it, and a procedure for passing that gradient on to its parents — the inputs it was built from.

The third one is the whole trick. An object knows only how it was made; it knows nothing about the graph as a whole. When a gradient arrives from upstream, it multiplies by its local derivative and hands the result to its parents. That is all it does. Wake these workers up one at a time starting from the output, and the gradient reaches every leaf. Nobody holds the big picture, yet the whole thing works — that is the pleasure of the autodiff design.

The skeleton: a Value that can add and multiply

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None   # how to pass gradients on (default: do nothing)
        self._prev = set(_children)     # the inputs this was built from
        self._op = _op                  # a label, for display only

    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

    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(ab)/da = b
            other.grad += self.data * out.grad   # d(ab)/db = a
        out._backward = _backward
        return out

__add__ and __mul__ are Python's operator overloading hooks. With them in place you write a + b like normal, and nodes quietly accumulate behind your back. This is why you never have to "declare" the computation graph: run the forward code and the graph falls out as a byproduct.

The other if isinstance(other, Value) else Value(other) at the top of each method is what lets raw numbers mix in, so x * 2 works. The constant becomes a node too and dutifully gets a gradient that nobody ever reads, which is harmless. Normalize the types at the door and every branch downstream disappears — that is how the implementation stays short.

Why store a function instead of a number

Storing local derivatives as plain numbers looks sufficient. For multiplication you could just attach the tuple (other.data, self.data) and be done. For simple operations that works.

It falls apart the moment you meet an operation whose local derivative comes from the output rather than the input. Add $\tanh$ and you see it.

Read it as: the slope of tanh is its own output, squared and subtracted from one. The point is that $x$ does not appear on the right-hand side. To know the slope you do not need the input — the output you already computed is enough. The exponential is the same story: $\frac{d}{dx}e^x = e^x$, so the output is the slope.

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,), 'tanh')
        def _backward():
            self.grad += (1 - t * t) * out.grad
        out._backward = _backward
        return out

Notice that _backward reaches both t and out. That is a Python closure: the function remembers the variables that existed when it was defined. With this shape, operations whose derivative depends on the input and operations whose derivative depends on the output both fit the same slot. PyTorch can ship hundreds of operations precisely because each one registers as a pair: how to compute forward, and how to hand gradients back.

Turn that around and adding a new operation is always the identical chore. Compute the forward value, build out, write a _backward that distributes the local derivative, attach it. ReLU, exp, log — the recipe never changes. Because the cost of adding one operation is constant, the library can keep growing its op set indefinitely. The _op label, incidentally, plays no part in the math; it exists so you can draw the graph when debugging.

Drag the point on the curve below and you can confirm the claim directly: for tanh, the slope at a point is readable from its height alone.

Get the order wrong and it breaks silently