{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Step 2 — Wire a Neural Network by Hand\n\n**Solution notebook.** A one-hidden-layer ReLU network, built unit by\nunit and then as two matrix layers, that classifies points of the plane\nas inside or outside the diamond $|x_1|+|x_2|=1$ — *exactly*, with\nweights written down on paper rather than learned. No corpus, no\ntraining, no autograd: this step is about seeing every number in a\nforward pass. It runs in a few seconds.\n\nThe construction is Lecture 2, Example 4.6 and Theorem 7.1; the\nderivation below is the \"On paper first\" section of the task, written\nout.",
   "id": "cell-00"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## On paper first\n\n**1. Four units whose activations sum to $r=|x_1|+|x_2|$.** For any real\n$t$, exactly one of $t$ and $-t$ is positive (or both are zero), so\n$\\operatorname{ReLU}(t)+\\operatorname{ReLU}(-t)=\\max(t,0)+\\max(-t,0)=|t|$.\nApply it to each coordinate: the four hidden units\n\n$$\nh=\\bigl(\\operatorname{ReLU}(x_1),\\ \\operatorname{ReLU}(-x_1),\\\n\\operatorname{ReLU}(x_2),\\ \\operatorname{ReLU}(-x_2)\\bigr)^\\top\n$$\n\nhave $h_1+h_2=|x_1|$ and $h_3+h_4=|x_2|$, so $\\mathbf 1^\\top h=r$.\n\n**2. As a matrix layer.** Each unit is $\\operatorname{ReLU}$ of an affine\nfunction of $\\mathbf x$, so $a_1=W_1\\mathbf x+b_1$, $h=\\operatorname{ReLU}(a_1)$ with\n\n$$\nW_1=\\begin{pmatrix}1&0\\\\-1&0\\\\0&1\\\\0&-1\\end{pmatrix}\\in\\mathbb R^{4\\times2},\n\\qquad b_1=\\mathbf 0\\in\\mathbb R^4,\n\\qquad a_1,h\\in\\mathbb R^4 .\n$$\n\n**3. The output layer.** We want $z_{\\mathrm{out}}=\\gamma(r-1)$ and\n$z_{\\mathrm{in}}=\\gamma(1-r)$. Since $r=\\mathbf 1^\\top h$,\n\n$$\nW_2=\\gamma\\begin{pmatrix}1&1&1&1\\\\-1&-1&-1&-1\\end{pmatrix}\\in\\mathbb R^{2\\times4},\n\\qquad\nb_2=\\gamma\\begin{pmatrix}-1\\\\1\\end{pmatrix}\\in\\mathbb R^2,\n\\qquad z=W_2h+b_2 .\n$$\n\n**4. The probability and the boundary.** For two classes, softmax\nreduces to a sigmoid of the logit difference (Lecture 2, Aside 3.1a):\n$p(\\mathrm{in}\\mid\\mathbf x)=\\sigma(z_{\\mathrm{in}}-z_{\\mathrm{out}})\n=\\sigma\\bigl(2\\gamma(1-|x_1|-|x_2|)\\bigr)$. Since $\\sigma$ is increasing\nwith $\\sigma(0)=\\tfrac12$, the classifier says *inside* exactly when\n$|x_1|+|x_2|<1$: the boundary is the diamond with vertices $(\\pm1,0)$,\n$(0,\\pm1)$. The sign of $1-r$ does not depend on $\\gamma$, so $\\gamma$\nrescales the logits — how *confident* the network is — without moving\nthe boundary at all.",
   "id": "cell-01"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.1 Write the hidden units separately\n\nFive test points, one per row. Each column of `h_units` is one hidden\nunit evaluated on all five points at once.",
   "id": "cell-02"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import torch\n\nX = torch.tensor([\n    [ 0.0,  0.0],   # inside\n    [ 0.5,  0.25],  # inside\n    [ 1.0,  0.0],   # boundary\n    [ 0.8,  0.5],   # outside\n    [-0.4, -0.2],   # inside, negative coordinates\n])\nprint(X.shape)      # (5, 2): five points, two features per point\n\nh_units = torch.stack([\n    torch.relu( X[:, 0]),\n    torch.relu(-X[:, 0]),\n    torch.relu( X[:, 1]),\n    torch.relu(-X[:, 1]),\n], dim=1)\n\nprint(h_units)\nprint(h_units.shape)                 # (5, 4)\nr_units = h_units.sum(dim=1)\nprint(r_units)                       # |x1| + |x2| for each point",
   "id": "cell-03"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> The four columns are active on the four half-planes $x_1>0$, $x_1<0$,\n> $x_2>0$, $x_2<0$ respectively — each unit measures how far the point\n> sits into \"its\" half-plane and is silent elsewhere. For\n> $(-0.4,-0.2)$: unit 1 sees $-0.4$ and outputs $0$; unit 2 sees $+0.4$\n> and outputs $0.4$; unit 3 sees $-0.2$, outputs $0$; unit 4 outputs\n> $0.2$. Row: $(0,\\,0.4,\\,0,\\,0.2)$, sum $0.6=|{-0.4}|+|{-0.2}|$. The\n> five rows are the checkpoint matrix.",
   "id": "cell-04"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.2 Assemble the first layer as a matrix\n\nThe same four units as one weight matrix and one bias vector. The\nlecture writes one input as a column and computes $W_1\\mathbf x$; code\nstores a batch with one sample per row, so the same product is\n`X @ W1.T`: $(5\\times2)(2\\times4)+(4)=(5\\times4)$, the bias added to\nevery row by broadcasting.",
   "id": "cell-05"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "W1 = torch.tensor([\n    [ 1.0,  0.0],\n    [-1.0,  0.0],\n    [ 0.0,  1.0],\n    [ 0.0, -1.0],\n])\nb1 = torch.zeros(4)\n\na1 = X @ W1.T + b1\nh = torch.relu(a1)\n\nprint(a1.shape, h.shape)             # both (5, 4)\nassert torch.allclose(h, h_units)\nprint('matrix layer == four separate units ✓')",
   "id": "cell-06"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.3 Wire the output layer and softmax\n\n$(5\\times4)(4\\times2)+(2)=(5\\times2)$: two logits per point, columns\nordered *outside, inside*. Subtracting each row's largest logit before\nexponentiating changes nothing — softmax is invariant under adding a\nconstant to every logit in a row, since the constant factors out of\nnumerator and denominator alike — but it keeps `exp` from overflowing\nwhen $\\gamma$ is large.",
   "id": "cell-07"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "gamma = 4.0\nW2 = gamma * torch.tensor([\n    [ 1.0,  1.0,  1.0,  1.0],   # outside logit\n    [-1.0, -1.0, -1.0, -1.0],   # inside logit\n])\nb2 = gamma * torch.tensor([-1.0, 1.0])\n\nlogits = h @ W2.T + b2            # (5, 2)\nshifted = logits - logits.max(dim=1, keepdim=True).values\nweights = shifted.exp()\nprobs = weights / weights.sum(dim=1, keepdim=True)\n\nprint(logits)\nprint(probs)\nprint(probs.sum(dim=1))",
   "id": "cell-08"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> Inside probabilities at $\\gamma=4$: **0.9997, 0.8808, 0.5000, 0.0832,\n> 0.9608** — i.e. $\\sigma(8)$, $\\sigma(2)$, $\\sigma(0)$, $\\sigma(-2.4)$,\n> $\\sigma(3.2)$, since $2\\gamma(1-r)=8(1-r)$ and $r=0,\\,0.75,\\,1,\\,1.3,\\,0.6$.\n> On the boundary point $(1,0)$ the logits tie and both probabilities\n> are exactly $\\tfrac12$; `argmax` would return index 0 (outside) there,\n> a tie-breaking convention of the software, not a property of the\n> classifier. Every row sums to 1.",
   "id": "cell-09"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.4 Package one forward pass\n\nThe two layers as a function that returns every intermediate value.\nThen the two assertions the task asks for: the hidden layer sums to\n$|x_1|+|x_2|$, and the inside probability is\n$\\sigma(2\\gamma(1-|x_1|-|x_2|))$ — Theorem 7.1, checked numerically.",
   "id": "cell-10"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def diamond_net(X, gamma=4.0):\n    W1 = torch.tensor([\n        [ 1.0,  0.0],\n        [-1.0,  0.0],\n        [ 0.0,  1.0],\n        [ 0.0, -1.0],\n    ], dtype=X.dtype)\n    b1 = torch.zeros(4, dtype=X.dtype)\n\n    W2 = gamma * torch.tensor([\n        [ 1.0,  1.0,  1.0,  1.0],\n        [-1.0, -1.0, -1.0, -1.0],\n    ], dtype=X.dtype)\n    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)\n\n    a1 = X @ W1.T + b1\n    h = torch.relu(a1)\n    logits = h @ W2.T + b2\n\n    shifted = logits - logits.max(dim=1, keepdim=True).values\n    weights = shifted.exp()\n    probs = weights / weights.sum(dim=1, keepdim=True)\n    return a1, h, logits, probs\n\na1, h, logits, probs = diamond_net(X)\nassert a1.shape == (5, 4)\nassert h.shape == (5, 4)\nassert logits.shape == (5, 2)\nassert probs.shape == (5, 2)\n\n# the two assertions of our own\nr = X.abs().sum(dim=1)                                   # |x1| + |x2|\nassert torch.allclose(h.sum(dim=1), r)\nassert torch.allclose(probs[:, 1], torch.sigmoid(2 * 4.0 * (1 - r)))\nprint('hidden sum = |x1|+|x2| ✓   p(in) = sigmoid(2γ(1-|x1|-|x2|)) ✓')",
   "id": "cell-11"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.5 Draw the decision surface\n\nEvaluate the network on a $201\\times201$ grid and plot the inside\nprobability. `meshgrid` turns two coordinate axes into all coordinate\npairs; flattening and stacking makes the usual one-point-per-row batch;\n`reshape` puts the answers back on the grid for the plotting routine.",
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import matplotlib.pyplot as plt\n\naxis = torch.linspace(-1.6, 1.6, 201)\ngx, gy = torch.meshgrid(axis, axis, indexing='xy')\ngrid = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1)\n\n_, _, _, grid_probs = diamond_net(grid, gamma=4.0)\np_inside = grid_probs[:, 1].reshape(gx.shape)\n\nplt.figure(figsize=(6, 5))\ncontour_plot = plt.contourf(gx.numpy(), gy.numpy(), p_inside.numpy(),\n                            levels=30, cmap='Purples')\nplt.contour(gx.numpy(), gy.numpy(), p_inside.numpy(),\n            levels=[0.5], colors='black', linewidths=2)\nplt.scatter(X[:, 0], X[:, 1], c='red', s=30)\nplt.xlabel('x1')\nplt.ylabel('x2')\nplt.axis('equal')\nplt.colorbar(contour_plot, label='p(inside | x)')\nplt.show()",
   "id": "cell-13"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "The black $p=\\tfrac12$ contour is the diamond with vertices $(\\pm1,0)$,\n$(0,\\pm1)$. A single logistic-regression unit can only draw one line in\nthis plane; four hidden ReLUs measure distances from the two axes, and\ntheir sum bends one line into four.",
   "id": "cell-14"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.6 Change confidence without changing geometry\n\nThe same grid at $\\gamma\\in\\{0.25,1,4,20\\}$, one colour scale for all\nfour (`vmin=0, vmax=1`, so the panels are comparable).",
   "id": "cell-15"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "gammas = [0.25, 1.0, 4.0, 20.0]\nfig, axes = plt.subplots(1, 4, figsize=(16, 4), sharex=True, sharey=True)\nfor ax, g in zip(axes, gammas):\n    _, _, _, gp = diamond_net(grid, gamma=g)\n    pi = gp[:, 1].reshape(gx.shape).numpy()\n    im = ax.contourf(gx.numpy(), gy.numpy(), pi, levels=torch.linspace(0, 1, 31).numpy(),\n                     cmap='Purples', vmin=0, vmax=1)\n    ax.contour(gx.numpy(), gy.numpy(), pi, levels=[0.5], colors='black', linewidths=2)\n    ax.set_title(f'γ = {g}')\n    ax.set_aspect('equal')\nfig.colorbar(im, ax=axes, label='p(inside | x)', shrink=0.8)\nplt.show()\n\n# the boundary does not move: p = 1/2 on the diamond's vertices for every γ\nvertices = torch.tensor([[1., 0.], [0., 1.], [-1., 0.], [0., -1.]])\nfor g in gammas:\n    _, _, _, vp = diamond_net(vertices, gamma=g)\n    assert torch.allclose(vp[:, 1], torch.full((4,), 0.5))\nprint('p(inside) = 1/2 on all four vertices, for every γ ✓')",
   "id": "cell-16"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "1. **The diamond itself** — every point with $|x_1|+|x_2|=1$ — has\n   probability exactly $\\tfrac12$ in all four plots, because there\n   $z_{\\mathrm{in}}=z_{\\mathrm{out}}=0$ whatever $\\gamma$ is.\n2. **As $\\gamma$ grows**, probabilities away from the boundary saturate\n   toward $0$ and $1$: the purple band of uncertainty narrows, and at\n   $\\gamma=20$ the plot is essentially a two-colour picture of the\n   diamond's indicator function.\n3. **As $\\gamma\\to0$**, both logits go to $0$ and every point tends to\n   $p=\\tfrac12$: the network is still *correct* (the sign of\n   $z_{\\mathrm{in}}-z_{\\mathrm{out}}$ never changes) but maximally\n   unconfident, and the whole plane fades to the same mid-purple.\n4. Softmax at temperature $T$ is $\\operatorname{softmax}(z/T)$. Here\n   every logit is proportional to $\\gamma$, so\n   $\\operatorname{softmax}(z(\\gamma))=\\operatorname{softmax}(z(1)/T)$\n   with $T=1/\\gamma$: changing $\\gamma$ **is** changing the temperature.\n   Temperature rescales a distribution's sharpness and leaves its argmax\n   — the geometry — untouched. (Lecture 7 returns to this when we\n   sample.)",
   "id": "cell-17"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2.7 Move and stretch the diamond\n\nFor a centre $(a,b)$ and radii $r_1,r_2$ we want the hidden activations\nto sum to $\\dfrac{|x_1-a|}{r_1}+\\dfrac{|x_2-b|}{r_2}$. Since\n$|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)$,\nthe four units are affine in $\\mathbf x$ with\n\n$$\nW_1=\\begin{pmatrix}1/r_1&0\\\\-1/r_1&0\\\\0&1/r_2\\\\0&-1/r_2\\end{pmatrix},\n\\qquad\nb_1=\\begin{pmatrix}-a/r_1\\\\ a/r_1\\\\ -b/r_2\\\\ b/r_2\\end{pmatrix},\n$$\n\nand the output layer is unchanged: inside when the sum is below $1$. The\nboundary $\\frac{|x_1-a|}{r_1}+\\frac{|x_2-b|}{r_2}=1$ has vertices\n$(a\\pm r_1,b)$ and $(a,b\\pm r_2)$.",
   "id": "cell-18"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def stretched_diamond_net(X, a=0.5, b=-0.25, r1=1.2, r2=0.6, gamma=4.0):\n    W1 = torch.tensor([\n        [ 1/r1,   0.0],\n        [-1/r1,   0.0],\n        [  0.0,  1/r2],\n        [  0.0, -1/r2],\n    ], dtype=X.dtype)\n    b1 = torch.tensor([-a/r1, a/r1, -b/r2, b/r2], dtype=X.dtype)\n    W2 = gamma * torch.tensor([[1., 1., 1., 1.], [-1., -1., -1., -1.]], dtype=X.dtype)\n    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)\n\n    h = torch.relu(X @ W1.T + b1)\n    logits = h @ W2.T + b2\n    return torch.softmax(logits, dim=1)        # same as the by-hand softmax above\n\na, b, r1, r2 = 0.5, -0.25, 1.2, 0.6\naxis2 = torch.linspace(-1.5, 2.5, 201)\ngx2, gy2 = torch.meshgrid(axis2, axis2, indexing='xy')\ngrid2 = torch.stack([gx2.reshape(-1), gy2.reshape(-1)], dim=1)\np2 = stretched_diamond_net(grid2, a, b, r1, r2)[:, 1].reshape(gx2.shape)\n\nplt.figure(figsize=(6, 5))\nplt.contourf(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=30, cmap='Purples')\nplt.contour(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=[0.5], colors='black', linewidths=2)\npredicted = torch.tensor([[a + r1, b], [a - r1, b], [a, b + r2], [a, b - r2]])\nplt.scatter(predicted[:, 0], predicted[:, 1], c='red', s=40, zorder=3, label='predicted vertices')\nplt.legend(); plt.axis('equal'); plt.xlabel('x1'); plt.ylabel('x2')\nplt.show()\n\nassert torch.allclose(stretched_diamond_net(predicted, a, b, r1, r2)[:, 1], torch.full((4,), 0.5))\nprint('p(inside) = 1/2 exactly at the four predicted vertices ✓')",
   "id": "cell-19"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Going further — a square needs another composition\n\n$\\max(u,v)=\\operatorname{ReLU}(u-v)+v$, so $\\max(|x_1|,|x_2|)$ needs one\nmore ReLU layer *after* the absolute values are formed — a composition,\nnot wider first layer. With $h$ the four units above, $u=h_1+h_2$ and\n$v=h_3+h_4$:\n\n$$\na_2=W_{2}'\\,h=\\begin{pmatrix}1&1&-1&-1\\\\0&0&1&1\\end{pmatrix}h\n=\\begin{pmatrix}u-v\\\\ v\\end{pmatrix},\n\\qquad\nh_2=\\operatorname{ReLU}(a_2)=\\begin{pmatrix}\\operatorname{ReLU}(u-v)\\\\ v\\end{pmatrix}\n$$\n\n(the second unit passes $v\\ge0$ through unchanged), and\n$\\mathbf 1^\\top h_2=\\max(u,v)$. The output layer is the same rule as\nbefore on this sum. Dimensions: $W_1\\in\\mathbb R^{4\\times2}$,\n$W_2'\\in\\mathbb R^{2\\times4}$, $W_3\\in\\mathbb R^{2\\times2}$.",
   "id": "cell-20"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def square_net(X, gamma=4.0):\n    h = torch.relu(X @ W1.T + b1)                                 # (N, 4): |x1|, |x2| split into halves\n    W2p = torch.tensor([[1., 1., -1., -1.], [0., 0., 1., 1.]])    # (2, 4)\n    h2 = torch.relu(h @ W2p.T)                                    # (N, 2): (ReLU(u - v), v)\n    W3 = gamma * torch.tensor([[1., 1.], [-1., -1.]])             # (2, 2)\n    b3 = gamma * torch.tensor([-1.0, 1.0])\n    return torch.softmax(h2 @ W3.T + b3, dim=1)\n\nm = torch.maximum(X[:, 0].abs(), X[:, 1].abs())\nassert torch.allclose(square_net(X)[:, 1], torch.sigmoid(2 * 4.0 * (1 - m)))\n\np_sq = square_net(grid)[:, 1].reshape(gx.shape)\nplt.figure(figsize=(5, 5))\nplt.contourf(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=30, cmap='Purples')\nplt.contour(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=[0.5], colors='black', linewidths=2)\nplt.axis('equal'); plt.title('max(|x1|, |x2|) < 1: three layers, one square')\nplt.show()",
   "id": "cell-21"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**Count the wiring.** The diamond network has $4\\cdot2+4+2\\cdot4+2=22$\nscalar parameters, of which $4+0+8+2=14$ are nonzero: four in $W_1$\n(one per unit, choosing an axis and a sign), eight in $W_2$ (every\nhidden unit feeds both logits, with opposite signs), two in $b_2$ (the\nthreshold $r=1$). Every one of them has a job you can name — which is\nthe point of building a network by hand before letting gradient descent\nfill the numbers in.\n\n→ Continue with [Step 3](../project/step-3): differentiate this very\nnetwork — first with an autograd engine you write yourself, then with\nPyTorch's — and then train the first model of the course whose weights\nnobody wrote down.",
   "id": "cell-22"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}