Step 2 — Wire a Network by Hand, Then Train One

Solution notebook. Two constructions. Part A builds a one-hidden-layer ReLU network, unit by unit and then as two matrix layers, that classifies points of the plane as inside or outside the diamond $|x_1|+|x_2|=1$ — exactly, with weights written down on paper rather than learned. Part B then rebuilds Step 1's bigram model as a differentiable parametric family and trains it by gradient descent, with a gradient derived by hand — no autograd yet. Both parts run in under a minute.

Part A is Lecture 2, Example 4.6 and Theorem 7.1; Part B is Lecture 2's multinomial logistic regression (Proposition 3.4) applied to one-hot character features. The derivations below are the "On paper first" sections of the task, written out.

Part A — Wire a Diamond Classifier by Hand

On paper first

1. Four units whose activations sum to $r=|x_1|+|x_2|$. For any real $t$, exactly one of $t$ and $-t$ is positive (or both are zero), so $\operatorname{ReLU}(t)+\operatorname{ReLU}(-t)=\max(t,0)+\max(-t,0)=|t|$. Apply it to each coordinate: the four hidden units

$$ h=\bigl(\operatorname{ReLU}(x_1),\ \operatorname{ReLU}(-x_1),\ \operatorname{ReLU}(x_2),\ \operatorname{ReLU}(-x_2)\bigr)^\top $$

have $h_1+h_2=|x_1|$ and $h_3+h_4=|x_2|$, so $\mathbf 1^\top h=r$.

2. As a matrix layer. Each unit is $\operatorname{ReLU}$ of an affine function of $\mathbf x$, so $a_1=W_1\mathbf x+b_1$, $h=\operatorname{ReLU}(a_1)$ with

$$ W_1=\begin{pmatrix}1&0\\-1&0\\0&1\\0&-1\end{pmatrix}\in\mathbb R^{4\times2}, \qquad b_1=\mathbf 0\in\mathbb R^4, \qquad a_1,h\in\mathbb R^4 . $$

3. The output layer. We want $z_{\mathrm{out}}=\gamma(r-1)$ and $z_{\mathrm{in}}=\gamma(1-r)$. Since $r=\mathbf 1^\top h$,

$$ W_2=\gamma\begin{pmatrix}1&1&1&1\\-1&-1&-1&-1\end{pmatrix}\in\mathbb R^{2\times4}, \qquad b_2=\gamma\begin{pmatrix}-1\\1\end{pmatrix}\in\mathbb R^2, \qquad z=W_2h+b_2 . $$

4. The probability and the boundary. For two classes, softmax reduces to a sigmoid of the logit difference (Lecture 2, Aside 3.1a): $p(\mathrm{in}\mid\mathbf x)=\sigma(z_{\mathrm{in}}-z_{\mathrm{out}}) =\sigma\bigl(2\gamma(1-|x_1|-|x_2|)\bigr)$. Since $\sigma$ is increasing with $\sigma(0)=\tfrac12$, the classifier says inside exactly when $|x_1|+|x_2|<1$: the boundary is the diamond with vertices $(\pm1,0)$, $(0,\pm1)$. The sign of $1-r$ does not depend on $\gamma$, so $\gamma$ rescales the logits — how confident the network is — without moving the boundary at all.

2.1 Write the hidden units separately

Five test points, one per row. Each column of h_units is one hidden unit evaluated on all five points at once.

import torch

X = torch.tensor([
    [ 0.0,  0.0],   # inside
    [ 0.5,  0.25],  # inside
    [ 1.0,  0.0],   # boundary
    [ 0.8,  0.5],   # outside
    [-0.4, -0.2],   # inside, negative coordinates
])
print(X.shape)      # (5, 2): five points, two features per point

h_units = torch.stack([
    torch.relu( X[:, 0]),
    torch.relu(-X[:, 0]),
    torch.relu( X[:, 1]),
    torch.relu(-X[:, 1]),
], dim=1)

print(h_units)
print(h_units.shape)                 # (5, 4)
r_units = h_units.sum(dim=1)
print(r_units)                       # |x1| + |x2| for each point

torch.Size([5, 2])
tensor([[0.0000, -0.0000, 0.0000, -0.0000],
        [0.5000, 0.0000, 0.2500, 0.0000],
        [1.0000, 0.0000, 0.0000, -0.0000],
        [0.8000, 0.0000, 0.5000, 0.0000],
        [0.0000, 0.4000, 0.0000, 0.2000]])
torch.Size([5, 4])
tensor([0.0000, 0.7500, 1.0000, 1.3000, 0.6000])

The four columns are active on the four half-planes $x_1>0$, $x_1<0$, $x_2>0$, $x_2<0$ respectively — each unit measures how far the point sits into "its" half-plane and is silent elsewhere. For $(-0.4,-0.2)$: unit 1 sees $-0.4$ and outputs $0$; unit 2 sees $+0.4$ and outputs $0.4$; unit 3 sees $-0.2$, outputs $0$; unit 4 outputs $0.2$. Row: $(0,\,0.4,\,0,\,0.2)$, sum $0.6=|{-0.4}|+|{-0.2}|$. The five rows are the checkpoint matrix.

2.2 Assemble the first layer as a matrix

The same four units as one weight matrix and one bias vector. The lecture writes one input as a column and computes $W_1\mathbf x$; code stores a batch with one sample per row, so the same product is X @ W1.T: $(5\times2)(2\times4)+(4)=(5\times4)$, the bias added to every row by broadcasting.

W1 = torch.tensor([
    [ 1.0,  0.0],
    [-1.0,  0.0],
    [ 0.0,  1.0],
    [ 0.0, -1.0],
])
b1 = torch.zeros(4)

a1 = X @ W1.T + b1
h = torch.relu(a1)

print(a1.shape, h.shape)             # both (5, 4)
assert torch.allclose(h, h_units)
print('matrix layer == four separate units ✓')

torch.Size([5, 4]) torch.Size([5, 4])
matrix layer == four separate units ✓

2.3 Wire the output layer and softmax

$(5\times4)(4\times2)+(2)=(5\times2)$: two logits per point, columns ordered outside, inside. Subtracting each row's largest logit before exponentiating changes nothing — softmax is invariant under adding a constant to every logit in a row, since the constant factors out of numerator and denominator alike — but it keeps exp from overflowing when $\gamma$ is large.

gamma = 4.0
W2 = gamma * torch.tensor([
    [ 1.0,  1.0,  1.0,  1.0],   # outside logit
    [-1.0, -1.0, -1.0, -1.0],   # inside logit
])
b2 = gamma * torch.tensor([-1.0, 1.0])

logits = h @ W2.T + b2            # (5, 2)
shifted = logits - logits.max(dim=1, keepdim=True).values
weights = shifted.exp()
probs = weights / weights.sum(dim=1, keepdim=True)

print(logits)
print(probs)
print(probs.sum(dim=1))

tensor([[-4.0000,  4.0000],
        [-1.0000,  1.0000],
        [ 0.0000,  0.0000],
        [ 1.2000, -1.2000],
        [-1.6000,  1.6000]])
tensor([[3.3535e-04, 9.9966e-01],
        [1.1920e-01, 8.8080e-01],
        [5.0000e-01, 5.0000e-01],
        [9.1683e-01, 8.3173e-02],
        [3.9166e-02, 9.6083e-01]])
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

Inside probabilities at $\gamma=4$: 0.9997, 0.8808, 0.5000, 0.0832, 0.9608 — i.e. $\sigma(8)$, $\sigma(2)$, $\sigma(0)$, $\sigma(-2.4)$, $\sigma(3.2)$, since $2\gamma(1-r)=8(1-r)$ and $r=0,\,0.75,\,1,\,1.3,\,0.6$. On the boundary point $(1,0)$ the logits tie and both probabilities are exactly $\tfrac12$; argmax would return index 0 (outside) there, a tie-breaking convention of the software, not a property of the classifier. Every row sums to 1.

2.4 Package one forward pass

The two layers as a function that returns every intermediate value. Then the two assertions the task asks for: the hidden layer sums to $|x_1|+|x_2|$, and the inside probability is $\sigma(2\gamma(1-|x_1|-|x_2|))$ — Theorem 7.1, checked numerically.

def diamond_net(X, gamma=4.0):
    W1 = torch.tensor([
        [ 1.0,  0.0],
        [-1.0,  0.0],
        [ 0.0,  1.0],
        [ 0.0, -1.0],
    ], dtype=X.dtype)
    b1 = torch.zeros(4, dtype=X.dtype)

    W2 = gamma * torch.tensor([
        [ 1.0,  1.0,  1.0,  1.0],
        [-1.0, -1.0, -1.0, -1.0],
    ], dtype=X.dtype)
    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)

    a1 = X @ W1.T + b1
    h = torch.relu(a1)
    logits = h @ W2.T + b2

    shifted = logits - logits.max(dim=1, keepdim=True).values
    weights = shifted.exp()
    probs = weights / weights.sum(dim=1, keepdim=True)
    return a1, h, logits, probs

a1, h, logits, probs = diamond_net(X)
assert a1.shape == (5, 4)
assert h.shape == (5, 4)
assert logits.shape == (5, 2)
assert probs.shape == (5, 2)

# the two assertions of our own
r = X.abs().sum(dim=1)                                   # |x1| + |x2|
assert torch.allclose(h.sum(dim=1), r)
assert torch.allclose(probs[:, 1], torch.sigmoid(2 * 4.0 * (1 - r)))
print('hidden sum = |x1|+|x2| ✓   p(in) = sigmoid(2γ(1-|x1|-|x2|)) ✓')

hidden sum = |x1|+|x2| ✓   p(in) = sigmoid(2γ(1-|x1|-|x2|)) ✓

2.5 Draw the decision surface

Evaluate the network on a $201\times201$ grid and plot the inside probability. meshgrid turns two coordinate axes into all coordinate pairs; flattening and stacking makes the usual one-point-per-row batch; reshape puts the answers back on the grid for the plotting routine.

import matplotlib.pyplot as plt

axis = torch.linspace(-1.6, 1.6, 201)
gx, gy = torch.meshgrid(axis, axis, indexing='xy')
grid = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1)

_, _, _, grid_probs = diamond_net(grid, gamma=4.0)
p_inside = grid_probs[:, 1].reshape(gx.shape)

plt.figure(figsize=(6, 5))
contour_plot = plt.contourf(gx.numpy(), gy.numpy(), p_inside.numpy(),
                            levels=30, cmap='Purples')
plt.contour(gx.numpy(), gy.numpy(), p_inside.numpy(),
            levels=[0.5], colors='black', linewidths=2)
plt.scatter(X[:, 0], X[:, 1], c='red', s=30)
plt.xlabel('x1')
plt.ylabel('x2')
plt.axis('equal')
plt.colorbar(contour_plot, label='p(inside | x)')
plt.show()

The black $p=\tfrac12$ contour is the diamond with vertices $(\pm1,0)$, $(0,\pm1)$. A single logistic-regression unit can only draw one line in this plane; four hidden ReLUs measure distances from the two axes, and their sum bends one line into four.

Part B — From Counts to Parameters

Rebuild Step 1's bigram model as a differentiable parametric family and train it by gradient descent — with a gradient derived by hand, no autograd. The pivot of the course: from estimating a table to optimizing a function.

Setup: data and tokenizer (Step 1's solution, reproduced)

Every notebook in this series is self-contained: it re-creates what it needs from earlier steps in one compact cell, so you can run it top to bottom without opening the others. On Colab, uncomment the download.

# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
import torch

with open('input.txt') as f:
    text = f.read()
vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
ids = torch.tensor(encode(text), dtype=torch.long)

assert len(text) == 1_115_394 and V == 65
print(f'{len(text):,} characters, vocab {V}')

1,115,394 characters, vocab 65

2.6 Data as tensors

Every adjacent pair, as parallel input/target vectors.

xs, ys = ids[:-1], ids[1:]
n = len(xs)
ar = torch.arange(n)
print(f'{n:,} training pairs')

1,115,393 training pairs

2.7 Forward pass

Row $a$ of $W$ is the logit vector for input character $a$, so W[xs] gathers all $n$ logit rows at once. The max-subtraction is the shift-invariance of softmax (Lecture 2, Prop. 1.2) spent on numerical stability — prove in one line that it changes nothing, then try removing it with 10× larger initial weights and watch the infs.

def loss_and_p(W, X=None, Y=None):
    X, Y = (xs, ys) if X is None else (X, Y)
    idx = torch.arange(len(X))
    logits = W[X]
    logits = logits - logits.max(dim=1, keepdim=True).values
    p = logits.exp()
    p = p / p.sum(dim=1, keepdim=True)
    return -p[idx, Y].log().mean(), p

2.8 Backward pass — by hand

The derivation (do it on paper before reading the code): for one pair $(a,b)$, $\ell = -z_b + \log\sum_c e^{z_c}$, so $\partial\ell/\partial z_j = p_j - \mathbf 1_{j=b}$ — probabilities minus target. Averaged over the batch, each example's $p - y$ row is accumulated into row $x_t$ of $W$; that is exactly what index_add_ does.

def grad_fn(W):
    _, p = loss_and_p(W)
    dlogits = p.clone()
    dlogits[ar, ys] -= 1.0          # p - y
    dlogits /= n
    dW = torch.zeros_like(W)
    dW.index_add_(0, xs, dlogits)   # route each row back to its input char
    return dW

Validate before training — in float64

Central finite differences against the analytic gradient. The dtype is the whole game here. The loss is a mean over a million examples, so a single entry of $\nabla_W$ is $\sim 2\times10^{-5}$; the difference $\mathcal L(W+hE)-\mathcal L(W-hE)\approx 2hg \approx 4\times10^{-8}$ sits below float32's resolution $\varepsilon|\mathcal L| \approx 3\times10^{-7}$, and the check fails at relative error ≈ 5.9 on a perfectly correct gradient. In float64 the same check passes at $\sim 3\times10^{-7}$. If your finite-difference check ever fails, suspect the dtype before the mathematics. (We restrict to a 20k-pair subset so each loss evaluation is cheap.)

torch.manual_seed(1337)
sub = 20_000
Xs, Ys = xs[:sub], ys[:sub]
Wt = (torch.randn(V, V) * 0.1).double()

# analytic gradient on the subset
_, p = loss_and_p(Wt, Xs, Ys)
dl = p.clone(); dl[torch.arange(sub), Ys] -= 1.0; dl /= sub
g = torch.zeros_like(Wt); g.index_add_(0, Xs, dl)

h, errs = 1e-3, []
for (i, j) in [(5, 7), (20, 40), (33, 2), (60, 60), (12, 12)]:
    Wp = Wt.clone(); Wp[i, j] += h
    Wm = Wt.clone(); Wm[i, j] -= h
    num = (loss_and_p(Wp, Xs, Ys)[0] - loss_and_p(Wm, Xs, Ys)[0]) / (2 * h)
    errs.append(abs((num - g[i, j]) / g[i, j]).item())
print(f'max relative error (float64): {max(errs):.2e}')
assert max(errs) < 1e-5
print('gradient verified — never trust one you haven\'t finite-differenced')

max relative error (float64): 3.29e-07
gradient verified — never trust one you haven't finite-differenced

The count model, for comparison

Step 1's answer, which two theorems say gradient descent must approach: the MLE is the normalized count matrix (Lecture 1, Thm 3.3), and the loss is convex in $W$ (Lecture 2, Ex. 2), so GD cannot be trapped elsewhere.

N = torch.zeros((V, V), dtype=torch.long)
N.index_put_((xs, ys), torch.ones(n, dtype=torch.long), accumulate=True)
P = (N + 1).float()
P = P / P.sum(1, keepdim=True)
ce_count = -P[xs, ys].log().mean().item()
print(f'count model (add-one) CE: {ce_count:.4f}')

count model (add-one) CE: 2.4549

2.9 Gradient descent

First, a learning-rate sweep — 30 steps each from the same start. The reference numbers:

lr	0.1	1	50	500
loss @ 30	4.158	4.024	2.742	5.401

The tiny rates are converging, uselessly slowly; 500 has gone above its starting point (overshooting the valley every step). Why does this convex problem tolerate lr 50 when Step 3's network will want 0.1? Because the gradient here is a mean over $10^6$ examples spread across 4,225 entries — of order $10^{-5}$ — and the step size compensates. Learning rates are not transferable constants; they are inverse to the gradient scale of the specific problem.

torch.manual_seed(0)
W0 = torch.randn(V, V) * 0.01
for lr in (0.1, 1.0, 50.0, 500.0):
    W = W0.clone()
    for _ in range(30):
        W -= lr * grad_fn(W)
    print(f'lr={lr:>6}: loss after 30 steps = {loss_and_p(W)[0].item():.4f}')

lr=   0.1: loss after 30 steps = 4.1577

lr=   1.0: loss after 30 steps = 4.0238

lr=  50.0: loss after 30 steps = 2.7420

lr= 500.0: loss after 30 steps = 5.4008

Now the real run: 300 full-batch steps at lr 50. Reference trajectory: 4.1731 → 2.6530 → 2.5690 → 2.5352 → 2.5168 → 2.5052 → 2.4973 (loss 2.4971 after the final update).

W = W0.clone()
for step in range(301):
    L, _ = loss_and_p(W)
    if step % 50 == 0:
        print(f'step {step:>4}  loss {L.item():.4f}', flush=True)
    W -= 50.0 * grad_fn(W)

step    0  loss 4.1731

step   50  loss 2.6530

step  100  loss 2.5690

step  150  loss 2.5352

step  200  loss 2.5168

step  250  loss 2.5052

step  300  loss 2.4973

2.10 The punchline — read it carefully

Close, and visibly not there: 2.497 vs 2.455, still 0.042 nats apart after 300 steps. Both proofs are about the limit; the gap is about the rate, and the rate is governed by conditioning. Row $a$ of $W$ receives gradient mass proportional to how often character $a$ occurs — common characters get strong signal, rare ones almost none — so agreement should be excellent on busy rows and poor on rare ones. The next cell measures exactly that prediction.

Wsm = torch.softmax(W, dim=1)
diff = (Wsm - P).abs()
busy = N.sum(1) > 1000
print(f'final GD loss              : {loss_and_p(W)[0].item():.4f}')
print(f'count model CE             : {ce_count:.4f}')
print(f'max  |softmax(W) - P|      : {diff.max().item():.4f}')
print(f'mean |softmax(W) - P|      : {diff.mean().item():.5f}')
print(f'max  on rows with >1000 obs: {diff[busy].max().item():.4f}'
      f'   ({busy.sum().item()} such rows)')
worst = diff.max(dim=1).values.argmax().item()
print(f'worst row: {itos[worst]!r}, which occurs '
      f'{N.sum(1)[worst].item()} times in the corpus')

final GD loss              : 2.4971
count model CE             : 2.4549
max  |softmax(W) - P|      : 0.5421
mean |softmax(W) - P|      : 0.00522
max  on rows with >1000 obs: 0.0698   (53 such rows)
worst row: 'Q', which occurs 231 times in the corpus

Reference values: max 0.542 over all entries, 0.0052 on average, 0.070 on the 53 well-populated rows — and the worst row belongs to Q (231 occurrences against e's 94,611), exactly as the conditioning argument predicts: the rows GD leaves wrong are the rows the data barely constrains. Run the sampler from Step 1 on softmax(W) and the text is indistinguishable from the count model's, because sampling almost never visits the rows that are still wrong.

Which of Step 1's ingredients does $W \approx 0$ at initialization play the role of? Initializing all logits near zero means all rows start near uniform — the same place add-one smoothing shrinks toward — and 300 steps has not fully erased that prior from the starved rows. Early stopping is regularization; you have just watched it act.

→ Continue with Step 3: the last gradient anyone derives by hand, and the machine that derives all the rest.

Going further — change confidence without changing geometry

The same grid at $\gamma\in\{0.25,1,4,20\}$, one colour scale for all four (vmin=0, vmax=1, so the panels are comparable).

gammas = [0.25, 1.0, 4.0, 20.0]
fig, axes = plt.subplots(1, 4, figsize=(16, 4), sharex=True, sharey=True)
for ax, g in zip(axes, gammas):
    _, _, _, gp = diamond_net(grid, gamma=g)
    pi = gp[:, 1].reshape(gx.shape).numpy()
    im = ax.contourf(gx.numpy(), gy.numpy(), pi, levels=torch.linspace(0, 1, 31).numpy(),
                     cmap='Purples', vmin=0, vmax=1)
    ax.contour(gx.numpy(), gy.numpy(), pi, levels=[0.5], colors='black', linewidths=2)
    ax.set_title(f'γ = {g}')
    ax.set_aspect('equal')
fig.colorbar(im, ax=axes, label='p(inside | x)', shrink=0.8)
plt.show()

# the boundary does not move: p = 1/2 on the diamond's vertices for every γ
vertices = torch.tensor([[1., 0.], [0., 1.], [-1., 0.], [0., -1.]])
for g in gammas:
    _, _, _, vp = diamond_net(vertices, gamma=g)
    assert torch.allclose(vp[:, 1], torch.full((4,), 0.5))
print('p(inside) = 1/2 on all four vertices, for every γ ✓')

p(inside) = 1/2 on all four vertices, for every γ ✓

1. The diamond itself — every point with $|x_1|+|x_2|=1$ — has probability exactly $\tfrac12$ in all four plots, because there $z_{\mathrm{in}}=z_{\mathrm{out}}=0$ whatever $\gamma$ is.
2. As $\gamma$ grows, probabilities away from the boundary saturate toward $0$ and $1$: the purple band of uncertainty narrows, and at $\gamma=20$ the plot is essentially a two-colour picture of the diamond's indicator function.
3. As $\gamma\to0$, both logits go to $0$ and every point tends to $p=\tfrac12$: the network is still correct (the sign of $z_{\mathrm{in}}-z_{\mathrm{out}}$ never changes) but maximally unconfident, and the whole plane fades to the same mid-purple.
4. Softmax at temperature $T$ is $\operatorname{softmax}(z/T)$. Here every logit is proportional to $\gamma$, so $\operatorname{softmax}(z(\gamma))=\operatorname{softmax}(z(1)/T)$ with $T=1/\gamma$: changing $\gamma$ is changing the temperature. Temperature rescales a distribution's sharpness and leaves its argmax — the geometry — untouched. (Lecture 7 returns to this when we sample.)

Going further — move and stretch the diamond

For a centre $(a,b)$ and radii $r_1,r_2$ we want the hidden activations to sum to $\dfrac{|x_1-a|}{r_1}+\dfrac{|x_2-b|}{r_2}$. Since $|x_1-a|/r_1=\operatorname{ReLU}\!\bigl(\tfrac{x_1-a}{r_1}\bigr)+\operatorname{ReLU}\!\bigl(\tfrac{a-x_1}{r_1}\bigr)$, the four units are affine in $\mathbf x$ with

$$ W_1=\begin{pmatrix}1/r_1&0\\-1/r_1&0\\0&1/r_2\\0&-1/r_2\end{pmatrix}, \qquad b_1=\begin{pmatrix}-a/r_1\\ a/r_1\\ -b/r_2\\ b/r_2\end{pmatrix}, $$

and the output layer is unchanged: inside when the sum is below $1$. The boundary $\frac{|x_1-a|}{r_1}+\frac{|x_2-b|}{r_2}=1$ has vertices $(a\pm r_1,b)$ and $(a,b\pm r_2)$.

def stretched_diamond_net(X, a=0.5, b=-0.25, r1=1.2, r2=0.6, gamma=4.0):
    W1 = torch.tensor([
        [ 1/r1,   0.0],
        [-1/r1,   0.0],
        [  0.0,  1/r2],
        [  0.0, -1/r2],
    ], dtype=X.dtype)
    b1 = torch.tensor([-a/r1, a/r1, -b/r2, b/r2], dtype=X.dtype)
    W2 = gamma * torch.tensor([[1., 1., 1., 1.], [-1., -1., -1., -1.]], dtype=X.dtype)
    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)

    h = torch.relu(X @ W1.T + b1)
    logits = h @ W2.T + b2
    return torch.softmax(logits, dim=1)        # same as the by-hand softmax above

a, b, r1, r2 = 0.5, -0.25, 1.2, 0.6
axis2 = torch.linspace(-1.5, 2.5, 201)
gx2, gy2 = torch.meshgrid(axis2, axis2, indexing='xy')
grid2 = torch.stack([gx2.reshape(-1), gy2.reshape(-1)], dim=1)
p2 = stretched_diamond_net(grid2, a, b, r1, r2)[:, 1].reshape(gx2.shape)

plt.figure(figsize=(6, 5))
plt.contourf(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=30, cmap='Purples')
plt.contour(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=[0.5], colors='black', linewidths=2)
predicted = torch.tensor([[a + r1, b], [a - r1, b], [a, b + r2], [a, b - r2]])
plt.scatter(predicted[:, 0], predicted[:, 1], c='red', s=40, zorder=3, label='predicted vertices')
plt.legend(); plt.axis('equal'); plt.xlabel('x1'); plt.ylabel('x2')
plt.show()

assert torch.allclose(stretched_diamond_net(predicted, a, b, r1, r2)[:, 1], torch.full((4,), 0.5))
print('p(inside) = 1/2 exactly at the four predicted vertices ✓')

p(inside) = 1/2 exactly at the four predicted vertices ✓

Going further — a square needs another composition

$\max(u,v)=\operatorname{ReLU}(u-v)+v$, so $\max(|x_1|,|x_2|)$ needs one more ReLU layer after the absolute values are formed — a composition, not wider first layer. With $h$ the four units above, $u=h_1+h_2$ and $v=h_3+h_4$:

$$ a_2=W_{2}'\,h=\begin{pmatrix}1&1&-1&-1\\0&0&1&1\end{pmatrix}h =\begin{pmatrix}u-v\\ v\end{pmatrix}, \qquad h_2=\operatorname{ReLU}(a_2)=\begin{pmatrix}\operatorname{ReLU}(u-v)\\ v\end{pmatrix} $$

(the second unit passes $v\ge0$ through unchanged), and $\mathbf 1^\top h_2=\max(u,v)$. The output layer is the same rule as before on this sum. Dimensions: $W_1\in\mathbb R^{4\times2}$, $W_2'\in\mathbb R^{2\times4}$, $W_3\in\mathbb R^{2\times2}$.

def square_net(X, gamma=4.0):
    h = torch.relu(X @ W1.T + b1)                                 # (N, 4): |x1|, |x2| split into halves
    W2p = torch.tensor([[1., 1., -1., -1.], [0., 0., 1., 1.]])    # (2, 4)
    h2 = torch.relu(h @ W2p.T)                                    # (N, 2): (ReLU(u - v), v)
    W3 = gamma * torch.tensor([[1., 1.], [-1., -1.]])             # (2, 2)
    b3 = gamma * torch.tensor([-1.0, 1.0])
    return torch.softmax(h2 @ W3.T + b3, dim=1)

m = torch.maximum(X[:, 0].abs(), X[:, 1].abs())
assert torch.allclose(square_net(X)[:, 1], torch.sigmoid(2 * 4.0 * (1 - m)))

p_sq = square_net(grid)[:, 1].reshape(gx.shape)
plt.figure(figsize=(5, 5))
plt.contourf(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=30, cmap='Purples')
plt.contour(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=[0.5], colors='black', linewidths=2)
plt.axis('equal'); plt.title('max(|x1|, |x2|) < 1: three layers, one square')
plt.show()

Count the wiring. The diamond network has $4\cdot2+4+2\cdot4+2=22$ scalar parameters, of which $4+0+8+2=14$ are nonzero: four in $W_1$ (one per unit, choosing an axis and a sign), eight in $W_2$ (every hidden unit feeds both logits, with opposite signs), two in $b_2$ (the threshold $r=1$). Every one of them has a job you can name — which is the point of building a network by hand before letting gradient descent fill the numbers in.

→ Continue with Step 3: differentiate this very network — first with an autograd engine you write yourself, then with PyTorch's — and then train the first model of the course whose weights nobody wrote down.