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. Your Tensor has the same interface as PyTorch's: we mark parameters with requires_grad=True, call loss.backward() and read .grad. In Section 7 we check that your engine and PyTorch give the same parameters after every update.

The lab follows the reverse-mode part of the autodiff lecture. For each operation we need its output (the forward pass) and 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	4	★ to ★★
2	The backward method	5	★
3	One gradient-descent step of the classifier	6	★

Section 3 defines a Tensor that only stores values; Sections 4 and 5 add its operations and backward. The reduction that undoes broadcasting, the topological order of the graph, 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. Exercises are rated from 1 to 3 stars based on expected difficulty.

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 and 2026-11-06) 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 = \partial L/\partial L = 1$. This is the seed.
• $\bar r = 2r\,\bar L = 10$.
• $\bar w = x\,\bar r = 20$ and $\bar b = \bar r = 10$.

Each step uses only one operation and the adjoint of its output. Reverse mode does this for any program, with one 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].$$

From this definition:

• $\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. The Tensor class

Our engine wraps a NumPy array in a Tensor with the same attributes as PyTorch's (see the table of tensor attributes in PT01):

• data: the NumPy array.
• requires_grad: whether we want gradients with respect to this tensor. We set it to True for parameters. The result of an operation requires gradients if any of its inputs does.
• grad: where backward stores the gradient. It starts as None. As in PyTorch, it is filled only for leaves with requires_grad=True.

Two more fields record the graph. Together they play the role of PyTorch's grad_fn:

• parents: the tensors this one was computed from, in order. Leaves have none.
• vjp: a function that maps the adjoint of this tensor to a tuple of adjoint contributions, one per parent. Leaves have vjp=None.

name is a label for plots.

The class below only stores values. We add the operations in Sections 4 and 5, with assignments such as Tensor.__add__ = add. This is the operator-overloading idea from the lecture: when we write a + b with two Tensors, Python calls Tensor.__add__, which computes the forward value and also stores a vjp function for the backward pass.

class Tensor:
    def __init__(self, data, requires_grad=False, parents=(), vjp=None, name='leaf'):
        self.data = np.array(data, dtype=np.float64)  # Always a float64 copy.
        self.parents = tuple(parents)  # The same node may appear twice, as in x * x.
        self.requires_grad = requires_grad or any(p.requires_grad for p in self.parents)
        self.grad = None
        self.vjp = vjp
        self.name = name

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

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

w = Tensor([0.2, -0.1], requires_grad=True, name='w')
print(w)
print('w.grad:', w.grad)

Tensor(w, data=[ 0.2 -0.1], requires_grad=True)
w.grad: None

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 given sum_to_shape then reduces it (next section).

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.

Broadcasting, and the first operations of Tensor

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), given in the next cell, 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, it sums 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, it sums over that axis with keepdims=True. Broadcasting stretched it.

For example, (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.

The same cell adds +, *, -, @, .T and .sum() to Tensor, using your VJPs. Note that:

• + and * pass their results through sum_to_shape.
• The output of every operation stores its inputs as parents and a vjp function. That function remembers the forward values it needs (here a and b).
• - 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.

The check after the cell compares the VJPs stored by + and * with PyTorch on broadcast shapes.

def sum_to_shape(v, shape):
    """Undo broadcasting: sum v over the axes that broadcasting added or stretched."""
    v = np.asarray(v, dtype=np.float64)
    while v.ndim > len(shape):
        v = v.sum(axis=0)
    for axis, size in enumerate(shape):
        if size == 1 and v.shape[axis] != 1:
            v = v.sum(axis=axis, keepdims=True)
    if v.shape != tuple(shape):
        raise ValueError('Target is not a broadcast-compatible input shape')
    return v

def add(a, b):
    b = as_tensor(b)
    def vjp(v):
        a_bar, b_bar = add_vjp(a.data, b.data, v)
        return sum_to_shape(a_bar, a.data.shape), sum_to_shape(b_bar, b.data.shape)
    return Tensor(a.data + b.data, parents=(a, b), vjp=vjp, name='+')

def multiply(a, b):
    b = as_tensor(b)
    def vjp(v):
        a_bar, b_bar = multiply_vjp(a.data, b.data, v)
        return sum_to_shape(a_bar, a.data.shape), sum_to_shape(b_bar, b.data.shape)
    return Tensor(a.data * b.data, parents=(a, b), vjp=vjp, name='*')

def matmul(a, b):
    b = as_tensor(b)
    if a.data.ndim not in (1, 2) or b.data.ndim not in (1, 2):
        raise ValueError('matmul supports vectors and matrices only')
    return Tensor(a.data @ b.data, parents=(a, b),
                  vjp=lambda v: matmul_vjp(a.data, b.data, v), name='@')

def transpose(a):
    if a.data.ndim not in (1, 2):
        raise ValueError('transpose supports vectors and matrices only')
    return Tensor(a.data.T, parents=(a,), vjp=lambda v: (v.T,), name='transpose')

def total(a, 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, a.data.shape).copy(),)
    return Tensor(a.data.sum(axis=axis), parents=(a,), vjp=vjp, name='sum')

# Addition and multiplication commute, so 3 * u can reuse u * 3.
Tensor.__add__ = Tensor.__radd__ = add
Tensor.__mul__ = Tensor.__rmul__ = multiply
Tensor.__neg__ = lambda a: a * -1.
Tensor.__sub__ = lambda a, b: a + (-as_tensor(b))
Tensor.__rsub__ = lambda a, b: as_tensor(b) + (-a)
Tensor.__matmul__ = matmul
Tensor.T = property(transpose)
Tensor.sum = total

def check_broadcasting():
    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.data.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('Broadcast addition and multiplication match PyTorch.')

check_broadcasting()

Broadcast addition and multiplication match PyTorch.

1b. Nonlinear operations ★

The VJP of a nonlinear operation needs values from the forward pass. For $y = \tanh(x)$ it is

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

which needs $y$. The cell after your functions adds tanh to Tensor; it keeps y in its vjp function and passes it to tanh_vjp. 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 1b: tanh_vjp')

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

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

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

def tanh(a):
    y = np.tanh(a.data)  # Saved for backward: the VJP needs the output.
    return Tensor(y, parents=(a,), vjp=lambda v: (tanh_vjp(y, v),), name='tanh')

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

Tensor.tanh = tanh
Tensor.relu = relu

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 1b: tanh and relu passed, including relu at 0.')

check_nonlinear_vjps()

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

5. Activity 2: backward ★

In PyTorch, loss.backward() computes the gradient of loss with respect to every leaf with requires_grad=True, and stores it in that leaf's .grad. We now write the same method. It calls the local VJPs over the whole graph, and two details need care: nodes used more than once, and the order in which nodes are handled.

Shared nodes. 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 below).
    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.data.item():g}' if node.name == 'constant' and node.data.size == 1 else node.name
        ax.text(x, y, f'{label}\n{node.data.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.], requires_grad=True, 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()

Order. 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 with a bug in it. It keeps the adjoints of the intermediate nodes in a dictionary, starting from the seed $\bar L = 1$, and handles each node as soon as a depth-first search starting from $L$ discovers it. Predict x.grad after this pass, then run the cell.

def premature_backward(output):
    """Looks fine, but handles each node as soon as it is discovered."""
    adjoints = {output: np.array(1.)}
    for node in graph_nodes(output):
        if not node.requires_grad:
            continue
        if node.vjp is None:
            node.grad = adjoints[node]
            continue
        for parent, contribution in zip(node.parents, node.vjp(adjoints[node])):
            if parent.requires_grad:
                adjoints[parent] = adjoints.get(parent, 0.) + contribution

order_x = Tensor(2., requires_grad=True, 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_backward(order_loss)
print(f'Premature backward: x.grad = {order_x.grad:g}   (correct: 44)')

Discovery order: L, a, u, x, constant, b
Premature backward: x.grad = 12   (correct: 44)

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 constant 3 needs no adjoint: it does not require gradients.

Two functions are given in the next cell:

• topological_order(output) returns every node that requires gradients and is reachable from output, each exactly once, with parents before the nodes that use them. It runs a depth-first search from output with an explicit stack (recursion would stop at Python's depth limit of about 1000) and appends each node once all its parents are in the order.
• initial_adjoint(output, gradient) returns the seed and raises the same errors as PyTorch: the seed is 1 for a scalar output, while a nonscalar output needs an explicit gradient of its shape.

def topological_order(output):
    """Nodes that require gradients, reachable from output, parents first."""
    order, seen, stack = [], set(), [(output, False)]
    while stack:
        node, parents_done = stack.pop()
        if parents_done:
            order.append(node)  # All its parents are already in the order.
        elif node not in seen:
            seen.add(node)
            stack.append((node, True))
            stack.extend((parent, False) for parent in reversed(node.parents)
                         if parent.requires_grad)
    return order

def initial_adjoint(output, gradient):
    """The seed of the backward pass, checked as in PyTorch."""
    if not output.requires_grad:
        raise RuntimeError('The tensor does not require grad and has no graph')
    if gradient is None:
        if output.data.shape != ():
            raise RuntimeError('grad can be implicitly created only for scalar outputs')
        return np.array(1.)
    gradient = np.array(gradient, dtype=np.float64)
    if gradient.shape != output.data.shape:
        raise RuntimeError('The gradient must have the shape of the output')
    return gradient

print('Topological order:', ', '.join(node.name for node in topological_order(order_loss)))

Topological order: x, u, a, b, L

Write backward(self, gradient=None). The cell then adds it to Tensor, so we can call loss.backward().

1. Start the dictionary of adjoints with the seed: adjoints = {self: initial_adjoint(self, gradient)}.
2. Go through topological_order(self) in reverse. For each node, take its adjoint adjoints[node].
   • If the node has a vjp, call it on the adjoint. It returns one contribution per parent. Add each contribution to the adjoint of the matching parent, skipping parents that do not require gradients. adjoints.get(parent, 0.) gives 0 for a parent that has no adjoint yet. A node that uses the same parent twice, as in x * x, sends it two contributions.
   • If the node is a leaf (its vjp is None), add the adjoint to node.grad, or set node.grad to it if node.grad is None.

Your code can follow the premature pass above. It differs in the order of the nodes and in how leaves update .grad.

As in PyTorch, backward returns nothing and accumulates: calling it twice doubles .grad, so between updates we reset gradients with x.grad = None (PT01). Intermediate nodes do not keep their adjoint, as in PyTorch without retain_grad().

def backward(self, gradient=None):
    raise NotImplementedError('Complete Activity 2: backward')

Tensor.backward = backward

Instructor reference: one reverse pass

Repeated parent entries still produce two contributions, which is necessary for x * x. The shape check turns a wrong local VJP into a clear error instead of a silent broadcast. Leaves receive a copy of their adjoint, so later accumulation never changes an array returned earlier.

def backward(self, gradient=None):
    adjoints = {self: initial_adjoint(self, gradient)}
    for node in reversed(topological_order(self)):
        adjoint = adjoints[node]
        if node.vjp is None:
            node.grad = adjoint.copy() if node.grad is None else node.grad + adjoint
            continue
        contributions = node.vjp(adjoint)
        if len(contributions) != len(node.parents):
            raise ValueError('A VJP must return one contribution per parent')
        for parent, contribution in zip(node.parents, contributions):
            if not parent.requires_grad:
                continue
            contribution = np.asarray(contribution, dtype=np.float64)
            if contribution.shape != parent.data.shape:
                raise ValueError('A VJP returned the wrong shape')
            adjoints[parent] = adjoints.get(parent, 0.) + contribution

Tensor.backward = backward

def check_backward():
    order_x.grad = None  # The premature pass above wrote into it.
    order_loss.backward()
    close(order_x.grad, 44., 'x.grad in the Section 5 example: u must collect both contributions first')
    order_loss.backward()
    close(order_x.grad, 88., 'a second backward call must add to .grad, as in PyTorch')
    assert order_u.grad is None, 'only leaves store .grad'
    close(order_x.data, 2., 'backward must not change the data')

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

    data = Tensor([[1., 2.], [3., 4.]])
    weight = Tensor([[1., -1.], [0., 2.], [3., 1.]], requires_grad=True)
    (data @ weight.T).sum().backward()
    assert data.grad is None, 'a tensor without requires_grad gets no .grad'
    close(weight.grad, np.tile(data.data.sum(axis=0), (3, 1)), 'matrix product inside a graph')
    bias = Tensor([1., 2., 3.], requires_grad=True)
    seed = np.arange(1., 10.).reshape(3, 3)
    (Tensor(np.arange(9.).reshape(3, 3)) + bias).backward(seed)
    close(bias.grad, seed.sum(axis=0), 'broadcast bias inside a graph')
    m = Tensor(np.arange(6.).reshape(2, 3), requires_grad=True)
    m.T.backward(np.arange(6.).reshape(3, 2))
    close(m.grad, np.arange(6.).reshape(3, 2).T, 'transpose')
    cube = Tensor(np.ones((2, 3, 4)), requires_grad=True)
    cube.sum(axis=(0, 2)).backward(np.array([1., 2., 3.]))
    close(cube.grad, 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., requires_grad=True)
    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.data, parents=(previous, previous), vjp=counted, name='+')
    repeated.backward()
    close(leaf.grad, 4096., 'doubling twelve times')
    assert counts == {i: 1 for i in range(12)}, f'each VJP must be called exactly once, got {counts}'
    deep = Tensor(0., requires_grad=True)
    for i in range(1500):
        deep = Tensor(deep.data + 1, parents=(deep,), vjp=lambda v: (v,), name='increment')
    deep.backward()  # A chain of 1500 operations.
    print('Activity 2: order, shared nodes, seeds, reductions and accumulation passed.')

check_backward()

Activity 2: order, shared nodes, seeds, reductions and accumulation 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.data.ndim != 2 or logits.data.shape[0] == 0 or logits.data.shape[1] == 0:
        raise ValueError('Logits must be a nonempty (batch, classes) matrix')
    if targets.shape != (len(logits.data),) 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.data.shape[1]).any():
        raise ValueError('Target index out of range')
    shifted = logits.data - logits.data.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, parents=(logits,), vjp=lambda v: (v * local_gradient,), name='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.]], requires_grad=True)
extreme_y = np.array([1, 0])
extreme_loss = cross_entropy(extreme_z, extreme_y)
tz = torch.tensor(extreme_z.data, requires_grad=True)
tloss = F.cross_entropy(tz, torch.tensor(extreme_y))
tloss.backward()
close(extreme_loss.data, tloss.detach().numpy(), 'loss on large logits')
extreme_loss.backward()
close(extreme_z.grad, tz.grad.numpy(), 'loss VJP on large logits')
assert np.isfinite(extreme_loss.data)
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 W and b in new Tensors with requires_grad=True: they are the parameters. X is data, so wrap it without.
2. Compute the logits X @ W.T + b and the loss with cross_entropy.
3. Call loss.backward(). The gradients are now in the .grad of the two parameter tensors.
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. Checking only the loss would miss some errors: 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, with empty .grad, and a new forward graph from them. In PT02 we updated the same parameters in place and cleared their gradients; fresh leaves have the same effect.

def classifier_step(X, targets, W, b, lr):
    weight = Tensor(W, requires_grad=True, name='W')
    bias = Tensor(b, requires_grad=True, name='b')
    logits = Tensor(X, name='X') @ weight.T + bias
    loss = cross_entropy(logits, targets)
    loss.backward()
    dw, db = weight.grad, bias.grad
    return dict(logits=logits.data.copy(), loss=loss.data.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).data.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 implements the core of PyTorch's autograd. During the forward pass, each operation records a VJP function. backward walks the graph in reverse topological order, adds up the contributions along every edge and stores the result in the .grad of the leaves.

PyTorch also supports hundreds of operations, GPUs and many data types, and it manages memory: by default it frees saved values after backward, so a second backward on the same graph needs retain_graph=True. Our engine keeps everything, supports matrix products of vectors and matrices only, and stores gradients as plain NumPy arrays, so you cannot differentiate a gradient a second time.

A Python loop does not create a cycle in the graph, because each iteration creates new nodes.

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.

Matching PyTorch on these tests does not prove the engine correct for every graph or operation.

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. The MLP uses only operations that your Tensor already has.

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, requires_grad=True, 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)
    loss.backward()
    return loss.data.item(), {name: params[name] - lr * leaves[name].grad 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, wraps it in a new Tensor x with requires_grad=True, runs function, calls backward and returns x.grad. 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, requires_grad=True, name='input')
        output = as_tensor(function(x))
        if output.data.shape != ():
            raise ValueError('grad requires a scalar output')
        if not output.requires_grad:  # The output does not depend on x.
            return np.zeros_like(x.data)
        output.backward()
        return x.grad
    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 gradient, and arrange the results in an array of shape output.shape + input.shape. Since .grad accumulates, reset it before each call.

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, requires_grad=True)
    output = as_tensor(function(x))
    result = np.zeros((output.data.size,) + x.data.shape)
    if output.requires_grad:  # Otherwise the output does not depend on x.
        for row in range(output.data.size):
            seed = np.zeros_like(output.data)
            seed.reshape(-1)[row] = 1
            x.grad = None
            output.backward(seed)
            result[row] = x.grad
    return result.reshape(output.data.shape + x.data.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, requires_grad=True)
custom_y = Tensor(A) @ (custom_x * custom_x)
custom_y.backward(v)
custom_vjp = custom_x.grad
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.data, 'forward value')
print('JVP:', jvp.numpy(), '| VJP:', custom_vjp)

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

/Users/simonescardapane/Repositories/teaching-materials/.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.