57. Autograd from scratch: the engine inside PyTorch in 120 lines

PyTorch's central trick is automatic differentiation: you write the forward computation, and it computes every gradient for you. Building a tiny version yourself removes the mystery. This chapter follows the design of Andrej Karpathy's micrograd, and is the code version of AI Ch 15 (backpropagation) and the chain rule in the Maths guide.

# labs/tinygrad.py: scalar reverse-mode autodiff + a small neural network
import math
import random

class Value:
    """A number that remembers how it was computed, so gradients can flow back."""

    __slots__ = ("data", "grad", "_backward", "_prev", "op")

    def __init__(self, data: float, _prev: tuple["Value", ...] = (), op: str = ""):
        self.data = data
        self.grad = 0.0                   # d(output)/d(self), filled in by backward()
        self._backward = lambda: None     # how to push my grad to my inputs
        self._prev = _prev
        self.op = op

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

    def __pow__(self, k: float):
        out = Value(self.data ** k, (self,), f"**{k}")
        def _backward():
            self.grad += k * self.data ** (k - 1) * out.grad
        out._backward = _backward
        return out

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

    def backward(self) -> None:
        order, seen = [], set()
        def visit(v):                     # topological sort: children before parents
            if v not in seen:
                seen.add(v)
                for p in v._prev:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0                   # d(out)/d(out)
        for v in reversed(order):         # chain rule, from the output back to the inputs
            v._backward()

    __radd__ = __add__
    __rmul__ = __mul__
    def __neg__(self): return self * -1
    def __sub__(self, other): return self + (-other)
    def __truediv__(self, other): return self * other ** -1
    def __repr__(self): return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"

class Neuron:
    def __init__(self, n_in: int):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(n_in)]
        self.b = Value(0.0)
    def __call__(self, x):
        return sum((wi * xi for wi, xi in zip(self.w, x)), self.b).tanh()
    def parameters(self):
        return self.w + [self.b]

class MLP:
    def __init__(self, n_in: int, sizes: list[int]):
        dims = [n_in] + sizes
        self.layers = [[Neuron(dims[i]) for _ in range(dims[i + 1])] for i in range(len(sizes))]
    def __call__(self, x):
        for layer in self.layers:
            x = [n(x) for n in layer]
        return x[0] if len(x) == 1 else x
    def parameters(self):
        return [p for layer in self.layers for n in layer for p in n.parameters()]

if __name__ == "__main__":
    # 1. Check one gradient by hand: f = a*b + c  ->  df/da = b = -3
    a, b, c = Value(2.0), Value(-3.0), Value(10.0)
    f = a * b + c
    f.backward()
    print(f, a.grad, b.grad, c.grad)                 # data=4, grads -3, 2, 1

    # 2. Train a 2-8-8-1 network to separate two classes of points.
    random.seed(1)
    X = [[random.uniform(-1, 1), random.uniform(-1, 1)] for _ in range(60)]
    Y = [1.0 if x0 * x0 + x1 * x1 < 0.5 else -1.0 for x0, x1 in X]   # inside / outside a circle
    model = MLP(2, [8, 8, 1])
    for step in range(200):
        preds = [model(x) for x in X]
        loss = sum((p - y) ** 2 for p, y in zip(preds, Y)) / len(X)   # mean squared error
        for p in model.parameters():
            p.grad = 0.0                                              # zero the old gradients
        loss.backward()                                               # backprop through everything
        lr = 0.1 if step < 150 else 0.03
        for p in model.parameters():
            p.data -= lr * p.grad                                     # gradient descent step
        if step % 50 == 0 or step == 199:
            acc = sum((p.data > 0) == (y > 0) for p, y in zip(preds, Y)) / len(X)
            print(f"step {step:3}  loss {loss.data:.4f}  accuracy {acc:.0%}")
    print(len(model.parameters()), "parameters")

Value(data=4.0000, grad=1.0000) -3.0 2.0 1.0
step   0  loss 1.0...  accuracy ...
...
step 199  loss 0.0...  accuracy 9x%
113 parameters

What you just built, and how it maps to PyTorch:

The Python features from earlier Parts are doing the work: operator overloading (Ch 16) makes a * b + c build a graph, closures (Ch 5.4) store each local derivative, and __slots__ (Ch 26) keeps the thousands of Value objects small. It is also painfully slow: one Python object per number. PyTorch's answer is the same engine over tensors, so one autograd node covers a million multiply-adds running in C++ or on the GPU.