PT03 · Automatic differentiation from scratch (instructor version)

In PT01 and PT02, loss.backward() computed every gradient for us. In this lab you build a small version of that machinery in NumPy and use it to train the PT02 penguin classifier. By the end, your engine and PyTorch produce the same numbers, update after update.

The lab follows the reverse-mode part of the autodiff lecture. Every operation knows two things: how to compute its output (the forward pass) and how to compute its vector-Jacobian product, or VJP. A backward pass calls the VJPs of all operations in the right order.

Activity	What you write	Section
1	The VJPs of addition, multiplication, matrix product, tanh and relu, and the reduction that undoes broadcasting	4
2	The backward pass over a whole graph	5
3	One gradient-descent step of the classifier	6

The Tensor class, graph plots, data and loss are given. Optional sections at the end train an MLP with your engine and build gradients and Jacobians as functions.

How to work. Run the cells in order. Each activity is followed by a check cell: if your code is wrong, the check stops with a message telling you what to look at. A function that still raises NotImplementedError stops the notebook at its first use, so complete the activities in order.

We assume knowledge of the material covered in PT01 (tensors and autograd) and PT02 (the batched linear classifier).

This instructor version includes executable solutions in cells tagged reference. The generated student version removes them and all saved outputs. An unfinished activity raises an error; it never silently calls a solution.

[!note] Written by Codex (2026-10-07), revised by Claude (2026-10-07) Adapted from Simone Scardapane's autodiff Colab and the shared PT01/PT02 notebooks. Sources are listed at the end.

1. Setup

The engine we implement will wrap NumPy with a PyTorch-like API for the backward pass. Everything runs on the CPU in double precision (float64) to avoid numerical disagreements. close and expect_error are small helpers used by the check cells.

from io import BytesIO
from pathlib import Path
from urllib.request import urlopen
import hashlib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch
import torch
import torch.nn.functional as F
from sklearn.model_selection import train_test_split

np.set_printoptions(precision=4, suppress=True)
torch.set_num_threads(1)
print('NumPy', np.__version__, '| PyTorch', torch.__version__)

def close(actual, expected, what='', *, atol=1e-9, rtol=1e-9):
    """Check shape and values; `what` says which quantity failed."""
    actual, expected = np.asarray(actual), np.asarray(expected)
    assert actual.shape == expected.shape, f'{what}: expected shape {expected.shape}, got {actual.shape}'
    np.testing.assert_allclose(actual, expected, atol=atol, rtol=rtol, err_msg=what)

def expect_error(kind, function):
    try:
        function()
    except kind:
        return
    raise AssertionError(f'Expected {kind.__name__}')

NumPy 2.5.4 | PyTorch 2.14.1

2. From the chain rule to VJPs

Start with an example you can do by hand. Take $w=3$, $x=2$, $b=-1$ and

$$r = wx + b,\qquad L = r^2.$$

The forward pass computes $r = 5$, then $L = 25$. Now go backward. For every value $z$ in the computation, write $\bar z = \partial L/\partial z$: how much $L$ changes when $z$ changes by a small amount. We call $\bar z$ the adjoint of $z$. Adjoints are computed from the end:

• $\bar L = 1$. This is the seed: $L$ changes exactly as much as itself.
• $\bar r = 2r\,\bar L = 10$.
• $\bar w = x\,\bar r = 20$ and $\bar b = \bar r = 10$.

Each step uses one operation and the adjoint of that operation's output, and nothing else. Reverse mode does this for any program: it splits the chain rule into one small step per operation. The next cell checks the numbers with finite differences.

w0, x0, b0 = 3., 2., -1.
r0 = w0 * x0 + b0
L0 = r0 ** 2

# Backward, from the seed to the inputs.
L0_bar = 1.
r0_bar = 2 * r0 * L0_bar
w0_bar, b0_bar = x0 * r0_bar, r0_bar

eps = 1e-6
loss_of = lambda w, b: (w * x0 + b) ** 2
w0_fd = (loss_of(w0 + eps, b0) - loss_of(w0 - eps, b0)) / (2 * eps)
b0_fd = (loss_of(w0, b0 + eps) - loss_of(w0, b0 - eps)) / (2 * eps)
print(f'w_bar = {w0_bar}   finite difference: {w0_fd:.6f}')
print(f'b_bar = {b0_bar}   finite difference: {b0_fd:.6f}')
close(w0_bar, w0_fd, 'w_bar', atol=1e-5)
close(b0_bar, b0_fd, 'b_bar', atol=1e-5)

w_bar = 20.0   finite difference: 20.000000
b_bar = 10.0   finite difference: 10.000000

The same idea for vectors and matrices

We saw that the derivative of $f$ at $x$ can be interpreted as a linear map, the JVP $\partial f(x)[u]$. It takes a direction $u$ in the input space and returns the resulting change of the output. Reverse mode uses the map that goes the other way. The VJP $\partial f(x)^*[v]$, the adjoint of the JVP, takes a direction $v$ in the output space and returns a direction in the input space. It is defined by

$$\langle v,\ \partial f(x)[u]\rangle = \langle \partial f(x)^*[v],\ u\rangle \quad\text{for every } u.$$

Here $\langle\cdot,\cdot\rangle$ is the sum of elementwise products, for vectors and matrices alike (np.sum(a * b)).

In a backward pass, $v$ is the adjoint of the output, and the VJP returns the adjoint of the input:

$$\bar x = \partial f(x)^*[\bar y].$$

Two facts follow immediately:

• $\bar x$ always has the same shape as $x$. Checking shapes catches many bugs.
• For a scalar loss, seeding $\bar L = 1$ and applying VJPs all the way back gives the gradient $\nabla L$ with respect to every input in a single backward pass. This is why deep learning uses reverse mode.

An operation with two inputs has one VJP per input: for $z = f(a, b)$ we need both $\bar a$ and $\bar b$.

Other names. JAX and some textbooks call the adjoint $\bar x$ a cotangent and the VJP a pullback. They are the same objects.

Predict. The lecture's example is $y = f(W) = Wx$. Take $W \sim (3, 4)$ and $x \sim (4)$. What are the shapes of $\bar y$, $\bar W$ and $\bar x$? Use the definition of the VJP to find a formula for $\bar W$. Hint: write $\langle v, Ux\rangle$ as a sum over the entries of $U$.

Instructor answer: $\bar y \sim (3)$, $\bar W \sim (3, 4)$ and $\bar x \sim (4)$. Since $\langle v, Ux\rangle = \sum_{i,j} v_i U_{ij} x_j = \langle v x^\top, U\rangle$, we get $\bar W = \bar y\,x^\top$ (an outer product) and, in the same way, $\bar x = W^\top \bar y$. Students implement exactly these formulas in Activity 1a.

3. A Tensor that remembers how it was made

Our engine wraps a NumPy array in a Tensor. Each Tensor stores:

• value: the array. It is read-only; to change a parameter, you create a new Tensor.
• parents: the tensors it was computed from, in order.
• vjp: a function that maps the adjoint of this tensor to a tuple of adjoint contributions, one per parent. Leaves (inputs and parameters) have vjp=None.
• name: a label for plots.

This is the operator-overloading idea from the lecture. When you write a * b with two Tensors, Python calls Tensor.__mul__, which computes the forward value and also stores a vjp function for later. That function remembers the forward values it needs (here a and b).

Compare with the lecture slide "Unpacking the PyTorch tensor": value plays the role of the data, and parents together with vjp play the role of grad_fn. There is no requires_grad and no .grad buffer. Instead, our backward will return a dictionary with the adjoint of every node.

The class is given. The VJPs it calls (add_vjp, multiply_vjp, ...) are yours to write in Activity 1, so the operators already work forward but cannot go backward yet. While you read, notice:

• + and * pass their results through sum_to_shape, which undoes broadcasting (Activity 1b).
• - is built from + and *, so it needs no VJP of its own.
• The VJPs of transpose and sum are complete. A transpose sends $\bar y$ back transposed. A sum copies $\bar y$ to every entry it added up, because each entry enters the sum with weight 1.

class Tensor:
    def __init__(self, value, parents=(), vjp=None, name='leaf'):
        self.value = np.array(value, dtype=np.float64, copy=True)
        self.value.setflags(write=False)
        self.parents = tuple(parents)  # The same node may appear twice, as in x * x.
        self.vjp = vjp
        self.name = name

    def __repr__(self):
        return f'Tensor({self.name}, shape={self.value.shape}, value={self.value})'

    def __add__(self, other):
        other = as_tensor(other)
        def vjp(v):
            a_bar, b_bar = add_vjp(self.value, other.value, v)
            return sum_to_shape(a_bar, self.value.shape), sum_to_shape(b_bar, other.value.shape)
        return Tensor(self.value + other.value, (self, other), vjp, '+')

    def __mul__(self, other):
        other = as_tensor(other)
        def vjp(v):
            a_bar, b_bar = multiply_vjp(self.value, other.value, v)
            return sum_to_shape(a_bar, self.value.shape), sum_to_shape(b_bar, other.value.shape)
        return Tensor(self.value * other.value, (self, other), vjp, '*')

    # Addition and multiplication commute, so 3 * u can reuse u * 3.
    __radd__ = __add__
    __rmul__ = __mul__

    def __neg__(self):
        return self * -1.

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

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

    def __matmul__(self, other):
        other = as_tensor(other)
        if self.value.ndim not in (1, 2) or other.value.ndim not in (1, 2):
            raise ValueError('matmul supports vectors and matrices only')
        return Tensor(self.value @ other.value, (self, other),
                      lambda v: matmul_vjp(self.value, other.value, v), '@')

    def tanh(self):
        y = np.tanh(self.value)  # Saved for backward: the VJP needs the output.
        return Tensor(y, (self,), lambda v: (tanh_vjp(y, v),), 'tanh')

    def relu(self):
        x = self.value  # Saved for backward: the VJP needs the sign of the input.
        return Tensor(np.maximum(x, 0.), (self,), lambda v: (relu_vjp(x, v),), 'relu')

    @property
    def T(self):
        if self.value.ndim not in (1, 2):
            raise ValueError('transpose supports vectors and matrices only')
        return Tensor(self.value.T, (self,), lambda v: (v.T,), 'transpose')

    def sum(self, axis=None):
        def vjp(v):
            # Put the summed axes back with size 1, then copy v along them.
            v = v if axis is None else np.expand_dims(v, axis)
            return (np.broadcast_to(v, self.value.shape).copy(),)
        return Tensor(self.value.sum(axis=axis), (self,), vjp, 'sum')

def as_tensor(value):
    """Wrap numbers and arrays as constant leaves; leave Tensors unchanged."""
    return value if isinstance(value, Tensor) else Tensor(value, name='constant')

Every operation creates a new node, and a node is a Python object with its own identity. In u + u, both inputs of + are the same node u. So the graph has one node u with two edges leaving it. In backward, each edge carries a contribution, and both must be counted.

Predict. With u = x * x, how many nodes does the graph of L = (u + u).sum() have? Run the cell to check. The plot draws each node once and each edge as an arrow, with its shape under the name.

def graph_nodes(output):
    # Collects nodes for drawing. This order is NOT a valid backward schedule (see Section 5).
    nodes, seen, pending = [], set(), [output]
    while pending:
        node = pending.pop()
        if node in seen:
            continue
        seen.add(node)
        nodes.append(node)
        pending.extend(reversed(node.parents))
    return nodes

def draw_graph(output):
    nodes = graph_nodes(output)
    remaining, depth = set(nodes), {}
    while remaining:
        ready = [n for n in nodes if n in remaining and all(p in depth for p in n.parents)]
        if not ready:
            raise ValueError('The graph must be acyclic')
        for node in ready:
            depth[node] = 0 if not node.parents else 1 + max(depth[p] for p in node.parents)
            remaining.remove(node)
    consumers = {node: [n for n in nodes if any(p is node for p in n.parents)] for node in nodes}
    # Draw each leaf one column before its first consumer, to avoid long edges.
    for node in nodes:
        if not node.parents and consumers[node]:
            depth[node] = min(depth[c] for c in consumers[node]) - 1
    # Place columns from right to left, each node near the nodes it feeds, to limit crossings.
    positions = {}
    for level in sorted(set(depth.values()), reverse=True):
        group = [n for n in reversed(nodes) if depth[n] == level]
        group.sort(key=lambda n: np.mean([positions[c][1] for c in consumers[n]]) if consumers[n] else 0.)
        for i, node in enumerate(group):
            positions[node] = (level * 2.2, i - (len(group) - 1) / 2)
    fig, ax = plt.subplots(figsize=(max(7, 2.2 * (max(depth.values()) + 1)), 3.5))
    for node in nodes:
        for slot, parent in enumerate(node.parents):
            repeated = sum(p is parent for p in node.parents) > 1
            radius = (0.16 if slot == 0 else -0.16) if repeated else 0
            arrow = FancyArrowPatch(positions[parent], positions[node], arrowstyle='->',
                                    mutation_scale=15, shrinkA=32, shrinkB=32,
                                    connectionstyle=f'arc3,rad={radius}', color='#526378')
            ax.add_patch(arrow)
    for node, (x, y) in positions.items():
        label = f'{node.value.item():g}' if node.name == 'constant' and node.value.size == 1 else node.name
        ax.text(x, y, f'{label}\n{node.value.shape}', ha='center', va='center',
                bbox=dict(boxstyle='round,pad=0.5', fc='#eef4fa', ec='#365b80'))
    ax.set_xlim(-0.8, max(x for x, y in positions.values()) + 0.8)
    ax.set_ylim(min(y for x, y in positions.values()) - 0.8,
                max(y for x, y in positions.values()) + 0.8)
    ax.axis('off')
    fig.tight_layout()
    return fig

graph_x = Tensor([1., 2.], name='x')
graph_u = graph_x * graph_x
graph_loss = (graph_u + graph_u).sum()
assert len(graph_nodes(graph_loss)) == 4
draw_graph(graph_loss)
plt.show()

4. Activity 1: local VJPs

1a. Addition, multiplication and matrix products

Write the VJP of each operation. Each function receives the two forward inputs a, b and the output adjoint v. It returns a tuple (a_bar, b_bar).

• Addition, $z = a + b$. A change in $a$ passes unchanged to $z$, so $\bar a = v$ and $\bar b = v$.
• Elementwise multiplication, $z = a \odot b$. Then $\bar a = v \odot b$ and $\bar b = v \odot a$.
• Matrix product, $z = AB$. Then $\bar A = v B^\top$ and $\bar B = A^\top v$. When one side is a vector, as in $z = Ax$, the transpose becomes an outer product: $\bar A = v x^\top$ (np.outer(v, x)) and $\bar x = A^\top v$. Handle all four cases: vector-vector, matrix-vector, vector-matrix and matrix-matrix. Use a.ndim and b.ndim to tell them apart.

In add_vjp and multiply_vjp, write the formulas above directly, as if a and b had the same shape. When they were broadcast, your result has the output's shape; the Tensor class then reduces it with sum_to_shape (Activity 1b).

Before running the checks, verify on paper that every result has the shape of the corresponding input.

def add_vjp(a, b, v):
    raise NotImplementedError('Complete Activity 1a: add_vjp')

def multiply_vjp(a, b, v):
    raise NotImplementedError('Complete Activity 1a: multiply_vjp')

def matmul_vjp(a, b, v):
    raise NotImplementedError('Complete Activity 1a: matmul_vjp')

Instructor reference: local VJPs

Writing the four rank combinations of the matrix product separately keeps their shapes visible. The vector cases are where students most often swap the outer product.

def add_vjp(a, b, v):
    return v, v

def multiply_vjp(a, b, v):
    return v * b, v * a

def matmul_vjp(a, b, v):
    if a.ndim == b.ndim == 1:
        return v * b, v * a
    if a.ndim == 2 and b.ndim == 1:
        return np.outer(v, b), a.T @ v
    if a.ndim == 1 and b.ndim == 2:
        return v @ b.T, np.outer(a, v)
    if a.ndim == b.ndim == 2:
        return v @ b.T, a.T @ v
    raise ValueError('matmul supports vectors and matrices only')

def check_rule(rule, operation, a, b, label, gen):
    """Compare a VJP rule with PyTorch and with finite differences."""
    seed = gen.normal(size=np.shape(operation(a, b)))
    a_bar, b_bar = rule(a, b, seed)
    ta = torch.tensor(a, requires_grad=True)
    tb = torch.tensor(b, requires_grad=True)
    (operation(ta, tb) * torch.tensor(seed)).sum().backward()
    close(a_bar, ta.grad.numpy(), f'{label}: first result (a_bar)')
    close(b_bar, tb.grad.numpy(), f'{label}: second result (b_bar)')
    # The VJP definition <v, df[u]> = <df*[v], u>, with df[u] from finite differences.
    ua, ub = gen.normal(size=np.shape(a)), gen.normal(size=np.shape(b))
    eps = 1e-6
    jvp = (operation(a + eps * ua, b + eps * ub) - operation(a - eps * ua, b - eps * ub)) / (2 * eps)
    close(np.sum(seed * jvp), np.sum(a_bar * ua) + np.sum(b_bar * ub),
          f'{label}: VJP definition', atol=1e-7, rtol=1e-7)

def check_elementwise_vjps():
    a, b = np.array([[1., -2.], [3., 0.5]]), np.array([[2., 4.], [-1., 3.]])
    v = np.array([[1., 2.], [3., 4.]])
    close(add_vjp(a, b, v)[0], v, 'add_vjp: a_bar should equal v')
    close(add_vjp(a, b, v)[1], v, 'add_vjp: b_bar should equal v')
    close(multiply_vjp(a, b, v)[0], v * b, 'multiply_vjp: a_bar should be v * b')
    close(multiply_vjp(a, b, v)[1], v * a, 'multiply_vjp: b_bar should be v * a')
    gen = np.random.default_rng(11)
    for shape in [(), (3,), (2, 3)]:
        a, b = gen.normal(size=shape), gen.normal(size=shape)
        check_rule(add_vjp, lambda x, y: x + y, a, b, f'add_vjp with shape {shape}', gen)
        check_rule(multiply_vjp, lambda x, y: x * y, a, b, f'multiply_vjp with shape {shape}', gen)
    print('Activity 1a: addition and multiplication passed.')

def check_matmul_vjp():
    gen = np.random.default_rng(12)
    for a_shape, b_shape in [((4,), (4,)), ((3, 4), (4,)), ((4,), (4, 3)), ((2, 4), (4, 3))]:
        a, b = gen.normal(size=a_shape), gen.normal(size=b_shape)
        check_rule(matmul_vjp, lambda x, y: x @ y, a, b,
                   f'matmul_vjp for shapes {a_shape} @ {b_shape}', gen)
    print('Activity 1a: matrix products passed for all four vector/matrix cases.')

check_elementwise_vjps()
check_matmul_vjp()

Activity 1a: addition and multiplication passed.
Activity 1a: matrix products passed for all four vector/matrix cases.

1b. Undoing broadcasting

In PT02, a bias $b \sim (C)$ is added to the logits $XW^\top \sim (B, C)$. NumPy broadcasts $b$: each entry $b_c$ is reused in all $B$ rows. In backward, each reuse sends back its own contribution, and these contributions must be added. So $\bar b$ is $v$ summed over the batch axis, which gives the shape $(C)$ of $b$.

sum_to_shape(v, shape) does this in general. It receives v with the broadcast output shape and returns an array with the target shape:

1. While v has more axes than shape, sum over the first axis. Broadcasting adds missing axes at the front.
2. For every axis where the target size is 1 but v is larger, sum over that axis with keepdims=True. Broadcasting stretched it.

Examples: (3, 3) to (3,) sums over axis 0; (2, 3, 4) to (1, 3, 1) sums over axes 0 and 2 and keeps them; any shape to () sums everything. Be careful with square batches: summing over the wrong axis still gives the right shape, only with the wrong values. The checks use non-symmetric arrays to catch this.

def sum_to_shape(v, shape):
    raise NotImplementedError('Complete Activity 1b: sum_to_shape')

def sum_to_shape(v, shape):
    v = np.asarray(v, dtype=np.float64)
    shape = tuple(shape)
    if len(shape) > v.ndim:
        raise ValueError('Target has more axes than the broadcast output')
    while v.ndim > len(shape):
        v = v.sum(axis=0)
    for axis, size in enumerate(shape):
        if size == 1:
            v = v.sum(axis=axis, keepdims=True)
        elif size != v.shape[axis]:
            raise ValueError('Target is not a broadcast-compatible input shape')
    return v.reshape(shape).copy()

def check_broadcasting():
    v = np.arange(1., 10.).reshape(3, 3)
    close(sum_to_shape(v, (3,)), [12., 15., 18.], 'sum_to_shape (3, 3) -> (3,): sum over axis 0')
    close(sum_to_shape(v, (1, 3)), [[12., 15., 18.]], 'sum_to_shape (3, 3) -> (1, 3)')
    close(sum_to_shape(v, (3, 1)), [[6.], [15.], [24.]], 'sum_to_shape (3, 3) -> (3, 1)')
    close(sum_to_shape(v, ()), 45., 'sum_to_shape (3, 3) -> (): sum everything')
    close(sum_to_shape(v, (3, 3)), v, 'sum_to_shape with the same shape: unchanged')
    close(sum_to_shape(np.ones((2, 3, 4)), (1, 3, 1)), np.full((1, 3, 1), 8.),
          'sum_to_shape (2, 3, 4) -> (1, 3, 1)')
    close(sum_to_shape(np.ones((2, 3, 4)), (4,)), np.full(4, 6.), 'sum_to_shape (2, 3, 4) -> (4,)')
    close(sum_to_shape(np.array(2.), ()), 2., 'sum_to_shape () -> ()')
    # Through the Tensor class, compared with PyTorch.
    gen = np.random.default_rng(5)
    for a_shape, b_shape in [((3, 3), (3,)), ((2, 3), ()), ((1, 3, 1), (2, 1, 4)), ((4, 1), (1, 3))]:
        a, b = gen.normal(size=a_shape), gen.normal(size=b_shape)
        for symbol, operation in [('+', lambda x, y: x + y), ('*', lambda x, y: x * y)]:
            output = operation(Tensor(a), Tensor(b))
            seed = gen.normal(size=output.value.shape)
            a_bar, b_bar = output.vjp(seed)
            ta = torch.tensor(a, requires_grad=True)
            tb = torch.tensor(b, requires_grad=True)
            (operation(ta, tb) * torch.tensor(seed)).sum().backward()
            close(a_bar, ta.grad.numpy(), f'{a_shape} {symbol} {b_shape}: a_bar')
            close(b_bar, tb.grad.numpy(), f'{a_shape} {symbol} {b_shape}: b_bar')
    print('Activity 1b: sum_to_shape and broadcast addition/multiplication passed.')

check_broadcasting()

Activity 1b: sum_to_shape and broadcast addition/multiplication passed.

1c. Nonlinear operations

Nonlinear operations show why reverse mode has to save values from the forward pass. The VJP of $y = \tanh(x)$ is

$$\bar x = (1 - y^2) \odot v,$$

which needs $y$ from the forward pass. Tensor.tanh keeps y in its vjp function and passes it to yours. Elementwise operations act on each entry separately, so their VJP is an elementwise product. As in the lecture's activation-function example, no Jacobian matrix is needed.

Write:

• tanh_vjp(y, v), where y is the forward output;
• relu_vjp(x, v), where x is the forward input. ReLU lets $v$ through where $x > 0$ and blocks it elsewhere. At exactly $x = 0$ ReLU has no derivative; use 0 there, as PyTorch does.

Each function returns one array, with the shape of v.

def tanh_vjp(y, v):
    raise NotImplementedError('Complete Activity 1c: tanh_vjp')

def relu_vjp(x, v):
    raise NotImplementedError('Complete Activity 1c: relu_vjp')

def tanh_vjp(y, v):
    return (1 - y ** 2) * v

def relu_vjp(x, v):
    return (x > 0) * v

def check_nonlinear_vjps():
    gen = np.random.default_rng(3)
    inputs = [np.array([-2., -0.5, 0., 0.5, 2.]), gen.normal(size=(2, 3))]
    for x in inputs:
        v = np.arange(1., x.size + 1).reshape(x.shape)
        for name, method, torch_function in [('tanh_vjp', Tensor.tanh, torch.tanh),
                                             ('relu_vjp', Tensor.relu, torch.relu)]:
            (x_bar,) = method(Tensor(x)).vjp(v)
            tx = torch.tensor(x, requires_grad=True)
            (torch_function(tx) * torch.tensor(v)).sum().backward()
            close(x_bar, tx.grad.numpy(), f'{name} at x = {x.round(3).tolist()}')
    print('Activity 1c: tanh and relu passed, including relu at 0.')

check_nonlinear_vjps()

Activity 1c: tanh and relu passed, including relu at 0.

5. Activity 2: the backward pass

Now we combine the local VJPs over a whole graph, and the order matters. Take a scalar $x = 2$ and

$$u = x^2,\qquad a = 3u,\qquad b = u^2,\qquad L = a + b.$$

By hand: $u$ feeds two branches, so its adjoint collects two contributions, $\bar u = 3 + 2u = 3 + 8 = 11$. Then $\bar x = 2x\,\bar u = 44$.

The cell below draws this graph and runs a backward pass that looks reasonable but is wrong. It handles each node as soon as a depth-first search starting from $L$ discovers it. initial_adjoints, defined at the top of the cell, creates a zero adjoint for every node and puts the seed on the output; you will reuse it. Predict the value of $\bar x$ this pass returns before running it.

def initial_adjoints(nodes, output, seed):
    """Zero adjoints for every node, and the seed for the output."""
    if seed is None:
        if output.value.shape != ():
            raise ValueError('A nonscalar output needs an explicit seed')
        seed = np.array(1.)
    seed = np.asarray(seed, dtype=np.float64)
    if seed.shape != output.value.shape:
        raise ValueError('The seed must have the shape of the output')
    adjoints = {node: np.zeros_like(node.value) for node in nodes}
    adjoints[output] = seed.copy()
    return adjoints

def premature_backward(output, seed=None):
    """Looks fine, but handles each node as soon as it is discovered."""
    nodes = graph_nodes(output)
    adjoints = initial_adjoints(nodes, output, seed)
    for node in nodes:
        if node.vjp is not None:
            for parent, contribution in zip(node.parents, node.vjp(adjoints[node])):
                adjoints[parent] += contribution
    return adjoints

order_x = Tensor(2., name='x')
order_u = order_x * order_x
order_u.name = 'u'
order_a = 3 * order_u
order_a.name = 'a'
order_b = order_u * order_u
order_b.name = 'b'
order_loss = order_a + order_b
order_loss.name = 'L'
draw_graph(order_loss)
plt.show()

print('Discovery order:', ', '.join(node.name for node in graph_nodes(order_loss)))
premature = premature_backward(order_loss)
print(f'Premature backward: x_bar = {premature[order_x]:g}   (correct: 44)')
print(f'u_bar reaches {premature[order_u]:g} in the end, but only after u was already handled.')

Discovery order: L, a, u, x, constant, b
Premature backward: x_bar = 12   (correct: 44)
u_bar reaches 11 in the end, but only after u was already handled.

The premature pass handled $u$ when only the contribution from $a$ had arrived. It sent $2x \cdot 3 = 12$ to $x$ and moved on. The contribution from $b$ reached $u$ later, but by then nobody passed it on to $x$.

The fix is the rule from the lecture's fan-out slide: handle a node only after all the nodes that use it. A topological order lists every node after its parents, for example $x, u, a, b, L$. Reading it backward ($L, b, a, u, x$) puts every node after all of its consumers. Several valid orders may exist, and any of them works. The graph also contains a constant leaf for the 3. It gets an adjoint too, which we never use.

Write two functions.

topological_order(output) returns every node reachable from output, each exactly once, with parents before the nodes that use them. Use an explicit stack rather than recursion: Python stops recursion at a depth of about 1000, and the checks build a chain of 1500 operations. One standard approach:

• Start with the stack [(output, False)] and an empty seen set.
• Pop (node, False): if the node is not in seen, add it, push (node, True), then push (parent, False) for each parent.
• Pop (node, True): all its parents are already in the order, so append the node.

backward(output, seed=None):

1. Get the nodes with topological_order(output).
2. Call initial_adjoints(nodes, output, seed) from the cell above. The seed is 1 by default for a scalar output; a nonscalar output needs an explicit seed of the same shape.
3. Go through the nodes in reverse order. For each node with a vjp, call it on the node's adjoint, and add (+=) each returned contribution to the adjoint of the matching parent. A node that uses the same parent twice, as in x * x, sends it two contributions.
4. Return the dictionary of adjoints.

Every call starts from fresh zeros, so calling backward twice gives the same result. PyTorch instead accumulates into .grad (PT01). Both compute the same derivatives; they store them differently.

def topological_order(output):
    raise NotImplementedError('Complete Activity 2: topological_order')

def backward(output, seed=None):
    raise NotImplementedError('Complete Activity 2: backward')

Instructor reference: one reverse pass

Nodes are appended after their parents. Repeated parent entries still produce two contributions in backward, which is necessary for x * x. The shape check turns a wrong local VJP into a clear error instead of a silent broadcast.

def topological_order(output):
    order, seen, stack = [], set(), [(output, False)]
    while stack:
        node, parents_done = stack.pop()
        if parents_done:
            order.append(node)
        elif node not in seen:
            seen.add(node)
            stack.append((node, True))
            stack.extend((parent, False) for parent in reversed(node.parents))
    return order

def backward(output, seed=None):
    nodes = topological_order(output)
    adjoints = initial_adjoints(nodes, output, seed)
    for node in reversed(nodes):
        if node.vjp is None:
            continue
        contributions = node.vjp(adjoints[node])
        if len(contributions) != len(node.parents):
            raise ValueError('A VJP must return one contribution per parent')
        for parent, contribution in zip(node.parents, contributions):
            contribution = np.asarray(contribution, dtype=np.float64)
            if contribution.shape != parent.value.shape:
                raise ValueError('A VJP returned the wrong shape')
            adjoints[parent] += contribution
    return adjoints

def check_backward():
    positions = {node: i for i, node in enumerate(topological_order(order_loss))}
    assert len(positions) == len(graph_nodes(order_loss)), 'topological_order: every node exactly once'
    assert all(positions[parent] < positions[node] for node in positions for parent in node.parents), \
        'topological_order: every parent must come before the nodes that use it'
    first = backward(order_loss)
    close(first[order_u], 11., 'u_bar in the Section 5 example: both branches must contribute')
    close(first[order_x], 44., 'x_bar in the Section 5 example')
    second = backward(order_loss)
    close(second[order_x], first[order_x], 'a second backward call must give the same result')
    assert first[order_x] is not second[order_x], 'each backward call must create fresh adjoints'
    close(order_x.value, 2., 'backward must not change the values')

    x = Tensor([2., -3.])
    close(backward((x * x).sum())[x], [4., -6.], 'x * x: a parent used twice receives two contributions')
    y = x * x + x
    close(backward(y, np.array([2., -1.]))[x], [10., 5.], 'backward with an explicit seed')
    close(backward((x - y).sum())[x], [-4., 6.], 'subtraction: x - (x * x + x) has gradient -2x')
    expect_error(ValueError, lambda: backward(y))
    expect_error(ValueError, lambda: backward(y, 1.))
    expect_error(ValueError, lambda: backward(y, np.ones((2, 1))))
    expect_error(ValueError, lambda: backward(Tensor([1.])))
    leaf = Tensor(3.)
    close(backward(leaf)[leaf], 1., 'a leaf can itself be the output')

    bias = Tensor([1., 2., 3.])
    z = Tensor(np.arange(9.).reshape(3, 3)) + bias
    seed = np.arange(1., 10.).reshape(3, 3)
    close(backward(z, seed)[bias], seed.sum(axis=0), 'broadcast bias inside a graph')
    m = Tensor(np.arange(6.).reshape(2, 3))
    close(backward(m.T, np.arange(6.).reshape(3, 2))[m], np.arange(6.).reshape(3, 2).T, 'transpose')
    cube = Tensor(np.ones((2, 3, 4)))
    close(backward(cube.sum(axis=(0, 2)), np.array([1., 2., 3.]))[cube],
          np.broadcast_to(np.array([1., 2., 3.])[None, :, None], (2, 3, 4)), 'sum over two axes')

    # Twelve shared additions: 13 nodes, each VJP called once, not 4096 separate paths.
    leaf = Tensor(1.)
    repeated = leaf
    counts = {}
    for i in range(12):
        previous = repeated
        def counted(v, key=i):
            counts[key] = counts.get(key, 0) + 1
            return (v, v)
        repeated = Tensor(2 * previous.value, (previous, previous), counted, '+')
    close(backward(repeated)[leaf], 4096., 'doubling twelve times')
    assert counts == {i: 1 for i in range(12)}, f'each VJP must be called exactly once, got {counts}'
    assert len(topological_order(repeated)) == 13
    deep = Tensor(0.)
    for i in range(1500):
        deep = Tensor(deep.value + 1, (deep,), lambda v: (v,), 'increment')
    assert len(topological_order(deep)) == 1501, 'a chain of 1500 operations'
    backward(deep)
    print('Activity 2: order, shared nodes, seeds, reductions and fresh adjoints passed.')

check_backward()

Activity 2: order, shared nodes, seeds, reductions and fresh adjoints passed.

6. Activity 3: training the PT02 classifier

Back to PT02's model: logits $Z = XW^\top + b$, with $X \sim (B, D)$, $W \sim (C, D)$ and $b \sim (C)$. The next cell repeats PT02's data setup: four penguin measurements, the same train/validation/test split, and scaling computed on the training set only. It reads the bundled CSV, or downloads it if you opened the notebook alone in Colab. Nothing to do here.

DATA_URL = 'https://raw.githubusercontent.com/allisonhorst/palmerpenguins/main/inst/extdata/penguins.csv'
DATA_SHA256 = 'f204db2c753b0937caac3cb35258562c14f073e4bbc76be24b4c51ce22767a93'
data_path = Path('data/penguins.csv')
if data_path.exists():
    data_bytes = data_path.read_bytes()
else:
    data_bytes = urlopen(DATA_URL, timeout=30).read()
assert hashlib.sha256(data_bytes).hexdigest() == DATA_SHA256, 'Dataset changed; use the bundled CSV.'
features = ['bill_length_mm', 'bill_depth_mm', 'flipper_length_mm', 'body_mass_g']
classes = ['Adelie', 'Chinstrap', 'Gentoo']
penguins = pd.read_csv(BytesIO(data_bytes)).dropna(subset=features + ['species'])
# Same float32 preprocessing as PT02; the values are then promoted to float64.
raw_X = torch.tensor(penguins[features].to_numpy(), dtype=torch.float32)
all_y = pd.Categorical(penguins.species, categories=classes).codes.copy().astype(np.int64)
indices = np.arange(len(all_y))
train_idx, rest_idx = train_test_split(indices, test_size=0.4, stratify=all_y, random_state=7)
val_idx, test_idx = train_test_split(rest_idx, test_size=0.5, stratify=all_y[rest_idx], random_state=7)
mean = raw_X[train_idx].mean(dim=0)
scale = raw_X[train_idx].std(dim=0, correction=0).clamp_min(1e-8)
Xtrain, Xval, Xtest = [((raw_X[idx] - mean) / scale).double().numpy()
                       for idx in (train_idx, val_idx, test_idx)]
ytrain, yval, ytest = [all_y[idx] for idx in (train_idx, val_idx, test_idx)]
assert Xtrain.shape == (205, 4) and Xval.shape == (68, 4) and Xtest.shape == (69, 4)
assert not (set(train_idx) & set(val_idx) or set(train_idx) & set(test_idx) or set(val_idx) & set(test_idx))
print('Examples:', len(ytrain), 'training,', len(yval), 'validation,', len(ytest), 'test')

Examples: 205 training, 68 validation, 69 test

The loss is given as a single operation with its own VJP; you already implemented it in PT02. For a batch of $B$ examples with integer labels $y_i$:

$$L(Z) = \frac1B\sum_{i}\Big[\log\sum_c e^{Z_{ic}} - Z_{i,y_i}\Big], \qquad \bar Z = \bar L\,\frac{\operatorname{softmax}(Z) - Y}{B},$$

where $Y$ is the one-hot matrix of the labels and the softmax acts on each row. Note the $1/B$: the loss is a mean, so its gradient is divided by the batch size as well. As in PT02, subtracting each row's maximum keeps the exponentials from overflowing.

torch_reference performs the same update with PyTorch; the checks compare your results with it. The last lines test the loss on very large logits.

def cross_entropy(logits, targets):
    logits = as_tensor(logits)
    targets = np.asarray(targets)
    if logits.value.ndim != 2 or logits.value.shape[0] == 0 or logits.value.shape[1] == 0:
        raise ValueError('Logits must be a nonempty (batch, classes) matrix')
    if targets.shape != (len(logits.value),) or not np.issubdtype(targets.dtype, np.integer):
        raise ValueError('Targets must be a batch of integer class indices')
    if (targets < 0).any() or (targets >= logits.value.shape[1]).any():
        raise ValueError('Target index out of range')
    shifted = logits.value - logits.value.max(axis=1, keepdims=True)
    exp = np.exp(shifted)
    normalizer = exp.sum(axis=1, keepdims=True)
    rows = np.arange(len(targets))
    value = (np.log(normalizer[:, 0]) - shifted[rows, targets]).mean()
    local_gradient = exp / normalizer       # softmax(Z)
    local_gradient[rows, targets] -= 1      # - Y
    local_gradient /= len(targets)          # / B
    return Tensor(value, (logits,), lambda v: (v * local_gradient,), 'cross-entropy')

def torch_reference(X, targets, W, b, lr):
    tw = torch.tensor(W, dtype=torch.float64, requires_grad=True)
    tb = torch.tensor(b, dtype=torch.float64, requires_grad=True)
    z = torch.tensor(X, dtype=torch.float64) @ tw.T + tb
    loss = F.cross_entropy(z, torch.tensor(targets, dtype=torch.long))
    dw, db = torch.autograd.grad(loss, (tw, tb))
    return dict(logits=z.detach().numpy(), loss=loss.item(),
                dW=dw.numpy(), db=db.numpy(),
                W=tw.detach().numpy() - lr * dw.numpy(),
                b=tb.detach().numpy() - lr * db.numpy())

extreme_z = Tensor([[1000., -1000., 0.], [-1000., 1000., 0.]])
extreme_y = np.array([1, 0])
extreme_loss = cross_entropy(extreme_z, extreme_y)
tz = torch.tensor(extreme_z.value.copy(), requires_grad=True)
tloss = F.cross_entropy(tz, torch.tensor(extreme_y))
tloss.backward()
close(extreme_loss.value, tloss.detach().numpy(), 'loss on large logits')
close(backward(extreme_loss)[extreme_z], tz.grad.numpy(), 'loss VJP on large logits')
assert np.isfinite(extreme_loss.value)
print('Supplied loss: value and gradient on large logits match PyTorch.')

Supplied loss: value and gradient on large logits match PyTorch.

Write classifier_step(X, targets, W, b, lr), one step of gradient descent with your engine:

1. Wrap X, W and b in new Tensors. W and b are the leaves whose adjoints you need.
2. Compute the logits X @ W.T + b and the loss with cross_entropy.
3. Call backward(loss) and read the adjoints of the two parameter leaves: they are the gradients.
4. Return a dictionary with the logits, the loss (a Python float), both gradients, and both updated parameters W - lr * dW and b - lr * db. Use the keys shown in the cell.

Do not modify the arrays you receive; return new ones. The checks compare every entry of the dictionary with PyTorch, on a small batch and on penguin batches, for two consecutive steps. A correct loss alone is not enough: if you forget to update the bias, the first loss is still right.

def classifier_step(X, targets, W, b, lr):
    # Return dict(logits=<array>, loss=<float>, dW=<array>, db=<array>,
    #             W=<updated array>, b=<updated array>).
    raise NotImplementedError('Complete Activity 3: classifier_step')

Instructor reference: classifier step

The updated arrays are not inserted into the old graph. The next call builds new leaves and a new forward graph from them, as PyTorch does on every iteration.

def classifier_step(X, targets, W, b, lr):
    weight, bias = Tensor(W, name='W'), Tensor(b, name='b')
    logits = Tensor(X, name='X') @ weight.T + bias
    loss = cross_entropy(logits, targets)
    adjoints = backward(loss)
    dw, db = adjoints[weight], adjoints[bias]
    return dict(logits=logits.value.copy(), loss=loss.value.item(),
                dW=dw, db=db, W=W - lr * dw, b=b - lr * db)

def check_classifier_step():
    cases = [(np.array([[1., 2.], [-1., 3.], [2., -2.]]), np.array([0, 1, 2])),
             (Xtrain[:7], ytrain[:7]), (Xtrain[:1], ytrain[:1])]
    for X, y in cases:
        gen = np.random.default_rng(13)
        W, b = gen.normal(size=(3, X.shape[1])) * 0.1, gen.normal(size=3) * 0.1
        before_W, before_b = W.copy(), b.copy()
        actual = classifier_step(X, y, W, b, 0.1)
        expected = torch_reference(X, y, W, b, 0.1)
        assert set(actual) == set(expected), f'classifier_step must return the keys {sorted(expected)}'
        assert isinstance(actual['loss'], float), 'the loss must be a Python float (use .item())'
        for key in expected:
            close(actual[key], expected[key], f'classifier_step: {key}')
        close(W, before_W, 'classifier_step must not modify W in place')
        close(b, before_b, 'classifier_step must not modify b in place')
        # A second step, from the updated parameters.
        again = classifier_step(X, y, actual['W'], actual['b'], 0.1)
        again_expected = torch_reference(X, y, expected['W'], expected['b'], 0.1)
        for key in again_expected:
            close(again[key], again_expected[key], f'classifier_step, second update: {key}')
    print('Activity 3: logits, loss, both gradients and two successive updates match PyTorch.')

check_classifier_step()

Activity 3: logits, loss, both gradients and two successive updates match PyTorch.

7. Training with your engine

The loop is given: 60 full-batch updates with learning rate 0.1, starting from zero weights as in PT02. At every update it runs your classifier_step and the PyTorch version side by side and checks that the parameters agree. The test set is not used.

Predict before running: should the NumPy and PyTorch curves overlap exactly, almost, or not at all?

def compare_training(steps=60, lr=0.1):
    W, b = np.zeros((3, 4)), np.zeros(3)
    ref_W, ref_b = W.copy(), b.copy()
    history = {'custom train': [], 'PyTorch train': [], 'custom validation': [], 'PyTorch validation': []}
    for update in range(steps):
        custom = classifier_step(Xtrain, ytrain, W, b, lr)
        reference = torch_reference(Xtrain, ytrain, ref_W, ref_b, lr)
        history['custom train'].append(custom['loss'])
        history['PyTorch train'].append(reference['loss'])
        W, b = custom['W'], custom['b']
        ref_W, ref_b = reference['W'], reference['b']
        history['custom validation'].append(cross_entropy(Tensor(Xval @ W.T + b), yval).value.item())
        with torch.no_grad():
            ref_z = torch.tensor(Xval) @ torch.tensor(ref_W).T + torch.tensor(ref_b)
            history['PyTorch validation'].append(F.cross_entropy(ref_z, torch.tensor(yval)).item())
        close(W, ref_W, f'W after update {update + 1}')
        close(b, ref_b, f'b after update {update + 1}')
    return history, W, b

training_history, final_W, final_b = compare_training()
assert training_history['custom train'][-1] < training_history['custom train'][0]
for split in ('train', 'validation'):
    close(training_history[f'custom {split}'], training_history[f'PyTorch {split}'], f'{split} losses')
fig, ax = plt.subplots(figsize=(7, 4))
for split, color in [('train', '#2463a6'), ('validation', '#a34b19')]:
    ax.plot(training_history[f'custom {split}'], color=color, label=f'NumPy {split}')
    ax.plot(training_history[f'PyTorch {split}'], '--', color=color, label=f'PyTorch {split}')
ax.set(xlabel='Update (training loss measured before each update)', ylabel='Mean cross-entropy')
ax.legend()
fig.tight_layout()
plt.show()
validation_accuracy = ((Xval @ final_W.T + final_b).argmax(axis=1) == yval).mean()
print(f"Training loss: {training_history['custom train'][0]:.4f} -> {training_history['custom train'][-1]:.4f}")
print(f'Validation accuracy after 60 updates: {validation_accuracy:.3f}')

Training loss: 1.0986 -> 0.2137
Validation accuracy after 60 updates: 0.956

8. What you built

Your engine does what PyTorch's autograd does at its core. During the forward pass, each operation records a VJP function. The backward pass walks the graph in reverse topological order and adds up the contributions along every edge.

PyTorch does much more: hundreds of operations, GPUs, many data types, and control over memory (by default it frees saved values after backward). Our engine keeps everything, supports matrix products of vectors and matrices only, and its backward returns plain NumPy arrays, so you cannot differentiate a gradient a second time.

A Python loop does not create a cycle in the graph: each iteration creates new nodes, so the graph stays acyclic.

Discuss

1. Why does the bias need sum_to_shape, while W does not?
2. In u + u, why is u one node but two edges? What goes wrong if backward keeps only the last contribution to each parent?
3. All your local VJPs passed their checks in Activity 1. How could backward still return a wrong gradient?
4. Why does tanh save its output rather than its input?

Instructor answers:

1. Each $b_c$ is reused for all $B$ examples, so its adjoint is the sum of $B$ contributions. $W$ enters through a matrix product, whose VJP $\bar W = \bar Z^\top X$ already sums over the batch.
2. Node identity lets backward compute $\bar u$ once and propagate it once. The two edges carry two contributions that must be added; overwriting gives, for example, $x$ instead of $2x$ as the derivative of (x * x).
3. A node handled before all its consumers have contributed propagates an incomplete adjoint (12 instead of 44 in Section 5). Correct local rules are not enough; the schedule matters too.
4. $\tanh'(x) = 1 - \tanh(x)^2$ is cheap to compute from the output, which the forward pass already produced. Saving the input would also work but would require computing $\tanh$ again.

Agreement with PyTorch here supports this graph and operation set on the tested cases, not an unrestricted autodiff implementation.

Optional A: an MLP with your engine

Do this after PT02's Section 9. Replace the linear model with PT02's MLP, which has one hidden layer of 16 ReLU units:

$$H = \operatorname{relu}(XW_1^\top + b_1),\qquad Z = HW_2^\top + b_2.$$

Write mlp_step(X, targets, params, lr). params is a dictionary of arrays W1 ~ (16, 4), b1 ~ (16), W2 ~ (3, 16) and b2 ~ (3). Return the loss (a Python float) and a new dictionary with the updated parameters, under the same keys. Your engine needs no change: the MLP only uses operations you have already written.

The cell after the exercise starts both engines from the same weights as PT02's make_mlp() with seed 7, and trains for 200 updates at learning rate 0.1.

def mlp_step(X, targets, params, lr):
    # Return (loss as a float, dict with updated W1, b1, W2, b2).
    raise NotImplementedError('Optional exercise: mlp_step')

def mlp_step(X, targets, params, lr):
    leaves = {name: Tensor(value, name=name) for name, value in params.items()}
    hidden = (Tensor(X, name='X') @ leaves['W1'].T + leaves['b1']).relu()
    loss = cross_entropy(hidden @ leaves['W2'].T + leaves['b2'], targets)
    adjoints = backward(loss)
    return loss.value.item(), {name: params[name] - lr * adjoints[leaves[name]] for name in params}

torch.manual_seed(7)
torch_mlp = torch.nn.Sequential(torch.nn.Linear(4, 16), torch.nn.ReLU(), torch.nn.Linear(16, 3)).double()
params = {'W1': torch_mlp[0].weight, 'b1': torch_mlp[0].bias,
          'W2': torch_mlp[2].weight, 'b2': torch_mlp[2].bias}
params = {name: value.detach().numpy().copy() for name, value in params.items()}
optimizer = torch.optim.SGD(torch_mlp.parameters(), lr=0.1)
mlp_losses = {'NumPy': [], 'PyTorch': []}
for update in range(200):
    loss, params = mlp_step(Xtrain, ytrain, params, 0.1)
    optimizer.zero_grad()
    torch_loss = F.cross_entropy(torch_mlp(torch.tensor(Xtrain)), torch.tensor(ytrain))
    torch_loss.backward()
    optimizer.step()
    mlp_losses['NumPy'].append(loss)
    mlp_losses['PyTorch'].append(torch_loss.item())
close(mlp_losses['NumPy'], mlp_losses['PyTorch'], 'MLP training losses', atol=1e-8, rtol=1e-8)
for name, layer, attribute in [('W1', 0, 'weight'), ('b1', 0, 'bias'), ('W2', 2, 'weight'), ('b2', 2, 'bias')]:
    close(params[name], getattr(torch_mlp[layer], attribute).detach().numpy(),
          f'{name} after 200 updates', atol=1e-8, rtol=1e-8)

fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(mlp_losses['NumPy'], label='NumPy engine')
ax.plot(mlp_losses['PyTorch'], '--', label='PyTorch')
ax.set(xlabel='Update', ylabel='Training cross-entropy', title='MLP, 16 hidden units, lr=0.1')
ax.legend()
fig.tight_layout()
plt.show()
hidden = np.maximum(Xval @ params['W1'].T + params['b1'], 0)
mlp_accuracy = ((hidden @ params['W2'].T + params['b2']).argmax(axis=1) == yval).mean()
print(f"MLP training loss: {mlp_losses['NumPy'][0]:.4f} -> {mlp_losses['NumPy'][-1]:.4f}")
print(f'MLP validation accuracy: {mlp_accuracy:.3f}')

MLP training loss: 1.1324 -> 0.0588
MLP validation accuracy: 0.985

Optional B: gradients as functions

JAX and torch.func provide grad(f), which returns a new function that computes the gradient of f. Write it with your engine. grad(function) returns a function that takes a NumPy array x, wraps it in a new Tensor, runs function, calls backward and returns the adjoint of the input. Raise ValueError if function does not return a scalar. If the output does not depend on the input, return zeros.

def grad(function):
    raise NotImplementedError('Optional exercise: grad')

def grad(function):
    def evaluate(value):
        x = Tensor(value, name='input')
        output = as_tensor(function(x))
        if output.value.shape != ():
            raise ValueError('grad requires a scalar output')
        return backward(output).get(x, np.zeros_like(x.value)).copy()
    return evaluate

original = np.array([1., -2., 3.])
close(grad(lambda x: (x * x).sum())(original), 2 * original, 'grad of sum(x * x)')
close(grad(lambda x: x.tanh().sum())(original), 1 - np.tanh(original) ** 2, 'grad of sum(tanh(x))')
close(grad(lambda x: Tensor(7.))(original), np.zeros_like(original), 'grad of a constant')
close(original, np.array([1., -2., 3.]), 'grad must not modify its input')
expect_error(ValueError, lambda: grad(lambda x: x * x)(original))
print('Optional B: grad checks passed.')

Optional B: grad checks passed.

Optional C: a Jacobian from VJPs

One VJP with seed $v$ gives $\partial f(x)^*[v]$. With $v = e_j$, the $j$-th basis vector of the output space, it gives row $j$ of the Jacobian, as in the lecture's reverse-mode slides. Write jacobian(function, value): run the forward pass once, then call backward once per output entry with a one-hot seed, and arrange the results in an array of shape output.shape + input.shape.

How many backward passes does a function from $\mathbb R^{1000}$ to $\mathbb R$ need? And a function from $\mathbb R$ to $\mathbb R^{1000}$? Which mode would you use for the second one?

def jacobian(function, value):
    raise NotImplementedError('Optional exercise: jacobian')

def jacobian(function, value):
    x = Tensor(value)
    output = as_tensor(function(x))
    result = np.zeros((output.value.size,) + x.value.shape)
    for row in range(output.value.size):
        seed = np.zeros_like(output.value)
        seed.reshape(-1)[row] = 1
        result[row] = backward(output, seed).get(x, np.zeros_like(x.value))
    return result.reshape(output.value.shape + x.value.shape)

Instructor answer: one backward pass per output entry, so 1 pass for $\mathbb R^{1000}\to\mathbb R$ and 1000 for $\mathbb R\to\mathbb R^{1000}$. Forward mode needs one JVP per input entry, so a single pass for the second function. Training needs only the first case: one scalar loss, one seed.

close(jacobian(lambda x: x * x, np.array([1., 2.])), np.diag([2., 4.]), 'Jacobian of x * x')
close(jacobian(lambda x: (x * x).sum(), np.array([1., 2.])), np.array([2., 4.]), 'Jacobian of a scalar function')
close(jacobian(lambda x: Tensor([3., 4.]), np.array([1., 2.])), np.zeros((2, 2)), 'Jacobian of a constant')
A_small = np.array([[1., 2.], [3., 4.], [5., 6.]])
close(jacobian(lambda x: Tensor(A_small) @ x, np.array([1., -1.])), A_small, 'Jacobian of a linear map')
matrix_input = np.arange(6.).reshape(2, 3)
J = jacobian(lambda x: x * x, matrix_input)
assert J.shape == (2, 3, 2, 3)
close(J.reshape(6, 6), np.diag(2 * matrix_input.reshape(-1)), 'Jacobian for a matrix input')
print('Optional C: Jacobian checks passed.')

Optional C: Jacobian checks passed.

Optional D: JVP and VJP side by side

Your engine only does reverse mode. torch.func has both modes, so we can compare them on $f(x) = A(x \odot x)$ with $A \sim (3, 2)$:

• JVP: $\partial f(x)[u] = A(2x \odot u)$, a vector of shape $(3)$;
• VJP: $\partial f(x)^*[v] = 2x \odot (A^\top v)$, a vector of shape $(2)$.

Predict the two shapes, then run the cell. It computes both with PyTorch, the VJP with your engine too, and checks the identity $\langle v, \partial f(x)[u]\rangle = \langle \partial f(x)^*[v], u\rangle$. Nothing to implement here.

A = np.array([[1., 2.], [-2., 1.], [3., -1.]])
point = np.array([1., -2.])
u = np.array([0.5, 1.])
v = np.array([1., -1., 2.])
tA = torch.tensor(A)
function = lambda x: tA @ (x * x)
value_jvp, jvp = torch.func.jvp(function, (torch.tensor(point),), (torch.tensor(u),))
value_vjp, torch_vjp = torch.func.vjp(function, torch.tensor(point))
(vjp,) = torch_vjp(torch.tensor(v))
custom_x = Tensor(point)
custom_y = Tensor(A) @ (custom_x * custom_x)
custom_vjp = backward(custom_y, v)[custom_x]
close(jvp.numpy(), A @ (2 * point * u), 'JVP formula')
close(vjp.numpy(), custom_vjp, 'your VJP against torch.func.vjp')
close(v @ jvp.numpy(), custom_vjp @ u, 'the VJP definition')
close(value_jvp.numpy(), custom_y.value, 'forward value')
print('JVP:', jvp.numpy(), '| VJP:', custom_vjp)

JVP: [-7. -6.  7.] | VJP: [18.  4.]

/Users/simonescardapane/Repositories/teaching-materials/build/nb-venv/lib/python3.12/site-packages/torch/jit/_script.py:1491: FutureWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(

Sources

• Simone Scardapane, original autodiff Colab: starter and reference fragments.
• Automatic differentiation lecture slides: derivatives as linear maps, JVPs and VJPs, reverse mode with operator overloading.
• PT01 and PT02 notebooks: tensors, autograd, the classifier and the MLP.
• JAX Autodiff Cookbook: JVPs, VJPs and the "cotangent" terminology. JAX is not a dependency.
• PyTorch torch.func.jvp and torch.func.vjp.
• Penguins data: Horst, Hill and Gorman (2020), palmerpenguins, CC0. Original collection by Kristen Gorman and Palmer Station LTER. The source URL and SHA-256 are in the data cell.