Build an autograd engine from scratch

Backpropagation sounds like a wall of calculus. The useful version is smaller: each operation knows how to pass a gradient to its inputs, then the graph is walked backward. By the end of this guide you will have the whole engine working in about 100 lines of Python.

What you're building

“Autograd” is short for automatic gradient. An autograd engine implements backpropagation: it computes the gradient of some output (a loss) with respect to all the inputs that produced it (the weights). That single capability is the mathematical core of libraries like PyTorch and JAX. The production versions add tensors, GPU kernels, and a lot of engineering.

We build it scalar-valued, working on individual numbers like -4.0 and 2.0, not tensors. That is deliberate. Real libraries pack thousands of scalars into n-dimensional arrays so the hardware can crunch them in parallel, but the math is the same. Stripping tensors away lets you see the mechanism without much machinery around it.

What a derivative means here

Forget the symbolic rules you memorized in school for a minute. Nobody writing neural nets derives the derivative of the whole network on paper; it would be tens of thousands of terms. What you need here is the definition: a derivative measures how sensitively the output responds when you nudge an input by a tiny amount.

Take any function, bump the input by a tiny h, and measure how much the output moved, divided by h. That ratio is the slope at that point:

def f(x):
    return 3*x**2 - 4*x + 5

h = 0.0001
x = 3.0
slope = (f(x + h) - f(x)) / h
print(slope)   # ~14.0, how fast f is rising at x = 3

We do not use this numerical trick for training, because it is too slow and imprecise. But it is the mental model for everything that follows, and later it is how you will check that your real gradients are correct.

The Value object: numbers that remember

A network is a huge expression, so we need a data structure that records how each number was produced. That's the Value object. It wraps a single scalar in .data, and this is the key part: when you add or multiply two Values, the result remembers its children (the inputs) and the operation that made it.

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0.0            # gradient, filled in by backprop
        self._prev = set(_children)
        self._op = _op            # the op that produced this node, e.g. '+'

    def __add__(self, other):
        return Value(self.data + other.data, (self, other), '+')

    def __mul__(self, other):
        return Value(self.data * other.data, (self, other), '*')

Now a + b doesn't just return a number. It returns a new Value that knows its parents were a and b, joined by a +. Build up an expression and you've quietly built a graph: every node points back at the nodes that made it. That graph is what we'll walk backwards.

Backprop is the chain rule, one node at a time

Give every Value a .grad field: the derivative of the final output L with respect to that node. Backprop fills these in, starting at the end. The base case is trivial: the derivative of L with respect to itself is 1. Then you walk backwards, and at every node you do one tiny, local step.

Here is the detail that makes it work. Each operation only knows its own local derivative: how its output moves when its inputs move. A + node has a local derivative of 1 for each input. A * node's local derivative for one input is just the other input's value. The chain rule says: to get how a node affects L, multiply its local derivative by the gradient flowing into it from above. Do that at every node, accumulating backwards, and you have computed every gradient in the graph.

Two patterns fall out and they are worth memorizing. A + node is a pure router: because its local derivative is 1, it copies the incoming gradient to each input unchanged. A * node is a swapper: it sends each input the incoming gradient times the other input's value.

From one neuron to an automatic backward()

Doing this by hand is fine for four nodes and painful for forty. So we automate it. Each operation stores a little closure, _backward, that knows how to push its output's gradient into its inputs, exactly the local rules above. Then backward() calls them in the right order: a node can only run after everything that depends on it has run. That ordering is a topological sort, walked in reverse.

def __add__(self, other):
    out = Value(self.data + other.data, (self, other), '+')
    def _backward():
        self.grad  += 1.0 * out.grad      # plus routes the gradient
        other.grad += 1.0 * out.grad
    out._backward = _backward
    return out

def __mul__(self, other):
    out = Value(self.data * other.data, (self, other), '*')
    def _backward():
        self.grad  += other.data * out.grad   # times swaps the inputs
        other.grad += self.data  * out.grad
    out._backward = _backward
    return out

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                 # base case: dL/dL = 1
    for v in reversed(topo):        # walk the graph backwards
        v._backward()

With that in place, loss.backward() gives every weight in the graph a gradient, no matter how deep the graph is. The same local rules that work for one neuron still work when the graph gets much larger.

You choose the size of your Lego bricks

A neuron squashes its output through tanh, which needs exponentiation, something + and * alone can't express. You have a choice: break tanh down into exp, divide, and add (more atomic ops to write backward passes for), or implement tanh as a single node with its own local derivative 1 - tanh(x)². Both give identical gradients. The only rule is: you can build an operation at any level of abstraction, as long as you know its local derivative. Atoms or composites; the engine doesn't care.

This is PyTorch, minus the speed

Everything you built has a direct twin in PyTorch. Your Value is its tensor; .data, .grad, and .backward() all exist there and do the same job. The difference is that PyTorch runs the same math over n-dimensional tensors so thousands of these operations happen in parallel on a GPU. When you register a custom op in PyTorch, you give it a forward pass and a local backward pass, the same pattern used here.

Now build it

Reading it helps, but the details settle when you type the code. Open the autograd-engine project and implement Value step by step: __add__, __mul__, tanh, then the automatic backward(). The tests compare your gradients with the numerical estimate from the first section, so you can tell when the engine is right.

Once your gradients flow end to end, the next move is to put them to work: train a neural net on top of your engine. Build Neuron, Layer, and MLP, then write the gradient-descent loop that makes the loss fall.