Autograd, the heart of it

This chapter sits in second place even though you do not really use what you build until step 05. The reason is simple: layers, loss, optimizer and training loop are only worth anything if gradients run through them. Autograd is not one part among eight, it is the foundation the other seven stand on. Skip it and you build the rest on sand.

And it is shorter than its reputation suggests. What you are about to write is roughly forty lines.

The goal

Differentiation is bookkeeping, not magic. The trick has two parts, and both are unspectacular:

First: take notes. Every result of a computation remembers which inputs it came from and by which operation. A computation quietly becomes a graph - in our case it simply hangs off every tensor as _children.

Second: walk backwards. Every operation knows how to distribute an incoming gradient across its inputs. For + it passes through unchanged to both sides. For * each side gets it multiplied by the value of the other (the product rule). For @ it is two matrix products, each with one transpose.

That is all the chain rule is: reverse the order, then let each link pass on its share. Once you have typed this out, you will never see a black box behind loss.backward() again.

Zero magic in this step

This is where the rule of the course bites hardest: NumPy is allowed, but torch, jax, tinygrad and above all autograd are forbidden. That last library is not called the same as this chapter by accident - it would take over the one insight the whole course exists for. NumPy may multiply and transpose; you write the derivatives.

The exercise

meintorch/tensor.py gets extended. New are grad, _children, _backward and the method backward(). The three operators keep their forward computation from chapter 01 and each gain a backward step.

import numpy as np


def as_tensor(x):
    return x if isinstance(x, Tensor) else Tensor(x)


def unbroadcast(grad, shape):
    """Reduces a gradient back to the shape it was broadcast from.

    Pure plumbing, so it comes for free: if NumPy copied one row across a whole
    batch on the way forward, the gradients of that batch have to be summed back
    up on the way backward.
    """
    while grad.ndim > len(shape):
        grad = grad.sum(axis=0)
    for axis, size in enumerate(shape):
        if size == 1 and grad.shape[axis] != 1:
            grad = grad.sum(axis=axis, keepdims=True)
    return grad


class Tensor:
    def __init__(self, data, _children=(), _op=""):
        self.data = np.asarray(data, dtype=np.float64)
        self.grad = np.zeros_like(self.data)
        self._children = tuple(_children)
        self._op = _op
        self._backward = lambda: None

    @property
    def shape(self):
        return self.data.shape

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

    def zero_grad(self):
        self.grad = np.zeros_like(self.data)

    def __add__(self, other):
        other = as_tensor(other)
        out = Tensor(self.data + other.data, (self, other), "+")

        def _backward():
            # TODO(du): how much of out.grad goes to self, how much to other?
            raise NotImplementedError

        out._backward = _backward
        return out

    def __mul__(self, other):
        other = as_tensor(other)
        out = Tensor(self.data * other.data, (self, other), "*")

        def _backward():
            # TODO(du): product rule - the partner is the local factor.
            raise NotImplementedError

        out._backward = _backward
        return out

    def __matmul__(self, other):
        other = as_tensor(other)
        out = Tensor(self.data @ other.data, (self, other), "@")

        def _backward():
            # TODO(du): two matrix products, each with one transpose.
            raise NotImplementedError

        out._backward = _backward
        return out

    # Convenience, for free: all of it derived from + and *
    def __neg__(self):
        return self * -1.0

    def __sub__(self, other):
        return self + (-as_tensor(other))

    def __radd__(self, other):
        return self + other

    def __rmul__(self, other):
        return self * other

    def backward(self):
        # TODO(du): build the topological order, set the seed gradient, walk
        # backwards and call every _backward().
        raise NotImplementedError

Hint 1 - the direction

For the three _backward functions: out.grad is the question "how much does the final result change if this value changes?". Your job is to pass that question on to self and other - each multiplied by whatever the operation does locally.

For backward(): you may only call a _backward once out.grad is complete. That means every node has to come after everyone who depends on it. A depth first search over _children produces exactly the reverse order - so reverse it once.

Hint 2 - more concrete

• +: the gradient passes unchanged to both sides, so self.grad += unbroadcast(out.grad, self.data.shape), and the same for other.
• *: self gets other.data * out.grad, other gets self.data * out.grad. Send both through unbroadcast.
• @: self.grad += out.grad @ other.data.T and other.grad += self.data.T @ out.grad. If you are unsure which order is right: there is only one where the shapes work out.
• backward(): a list topo plus a set against double visits, recursively the children first, then append yourself. After that self.grad = np.ones_like(self.data) as the seed and for t in reversed(topo): t._backward().

    # in __add__
        def _backward():
            self.grad += unbroadcast(out.grad, self.data.shape)
            other.grad += unbroadcast(out.grad, other.data.shape)

    # in __mul__
        def _backward():
            self.grad += unbroadcast(other.data * out.grad, self.data.shape)
            other.grad += unbroadcast(self.data * out.grad, other.data.shape)

    # in __matmul__
        def _backward():
            self.grad += out.grad @ other.data.T
            other.grad += self.data.T @ out.grad

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

        def visit(t):
            if t in visited:
                return
            visited.add(t)
            for child in t._children:
                visit(child)
            topo.append(t)

        visit(self)
        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            t._backward()

In goal mode

The autonomous run builds the same graph and checks it harder than a single milestone test would: it compares the analytic gradients against a numerical derivative. The instruction:

Extend `meintorch/tensor.py` with reverse-mode autograd, without external
libraries (NumPy allowed, torch/jax/tinygrad/autograd forbidden).
Every tensor gains `grad`, its children and a `_backward` closure; `backward()`
builds the topological order, sets the seed gradient to one and walks backwards.
Gradients are accumulated (+=), never overwritten. Cover addition, elementwise
multiplication and the matrix product, including broadcasting.
Write `tests/test_02_autograd.py`: the derivative of x²+x at 3 has to come out
as 7 automatically. Additionally check every gradient against a central
numerical difference (eps 1e-6, tolerance 1e-5).
Run `python -m pytest` until everything is green.

The milestone test

tests/test_02_autograd.py. The first test is the milestone of the chapter: the derivative of x² + x is 2x + 1, so 7 at the point 3 - and nobody wrote that down anywhere. It falls out of the bookkeeping.

import numpy as np

from meintorch.tensor import Tensor


def test_derivative_of_x_squared_plus_x():
    x = Tensor(3.0)
    y = x * x + x
    y.backward()
    assert np.allclose(x.grad, 7.0)


def test_gradients_accumulate_instead_of_overwriting():
    x = Tensor(2.0)
    y = x + x
    y.backward()
    assert np.allclose(x.grad, 2.0)


def test_matmul_gradients_have_the_shape_of_their_inputs():
    a = Tensor([[1.0, 2.0], [3.0, 4.0]])
    b = Tensor([[5.0], [6.0]])
    out = a @ b
    out.backward()
    assert a.grad.shape == a.shape
    assert b.grad.shape == b.shape
    assert np.allclose(b.grad, [[4.0], [6.0]])


def test_zero_grad_cleans_up():
    x = Tensor(3.0)
    (x * x).backward()
    x.zero_grad()
    assert np.allclose(x.grad, 0.0)