{ "cells": [ { "cell_type": "markdown", "id": "scalar-autograd-simple-01", "metadata": {}, "source": [ "# Scalar autograd: PyTorch, from scratch, then a fused sigmoid\n", "\n", "We will start with one graph:\n", "\n", "$$\n", "m=wx,\\qquad a=m+b,\\qquad e=a-y,\\qquad L=e^2\n", "$$\n", "\n", "with $w=2$, $x=3$, $b=1$, and $y=10$.\n", "\n", "First we let PyTorch differentiate it. Then we build the smallest useful version of the same idea\n", "ourselves. The point is not to replace PyTorch—it is to see what `.backward()` does. A final optional\n", "example then compares one fused sigmoid operation with the same sigmoid expanded into atomic operations.\n", "\n", "
\n", " \"Complete\n", "
\n", "

\n", " Forward computes left → right. Backward begins at L.grad = 1 and travels right → left.\n", " The teal number in each box is that node's final ∂L/∂node. On a phone, scroll sideways.\n", "

" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-02", "metadata": {}, "source": [ "## 1 · The whole example in PyTorch\n", "\n", "Each number below is a scalar tensor. We set `requires_grad=True` because we want PyTorch to calculate\n", "its loss derivative. In ordinary training, the target `y` would normally be fixed; here we track it only\n", "so the output matches our complete paper calculation.\n", "\n", "PyTorch keeps `.grad` automatically for the four leaf tensors. `retain_grad()` asks it to keep gradients\n", "for the intermediate values too." ] }, { "cell_type": "code", "execution_count": 1, "id": "scalar-autograd-simple-03", "metadata": { "purpose": "Create literal scalar tensors and run the lecture forward graph" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Forward: m = 6.0 , a = 7.0 , e = -3.0 , L = 9.0\n" ] } ], "source": [ "import torch\n", "\n", "torch.set_default_dtype(torch.float64)\n", "\n", "w = torch.tensor(2.0, requires_grad=True)\n", "x = torch.tensor(3.0, requires_grad=True)\n", "b = torch.tensor(1.0, requires_grad=True)\n", "y = torch.tensor(10.0, requires_grad=True)\n", "\n", "# Forward pass\n", "m = w * x\n", "a = m + b\n", "e = a - y\n", "L = e ** 2\n", "\n", "for node in (m, a, e, L):\n", " node.retain_grad()\n", "\n", "print(\"Forward: m =\", m.item(), \", a =\", a.item(),\n", " \", e =\", e.item(), \", L =\", L.item())" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-04", "metadata": {}, "source": [ "Now the important line:" ] }, { "cell_type": "code", "execution_count": 2, "id": "scalar-autograd-simple-05", "metadata": { "purpose": "Call backward and print the four leaf gradients directly" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "w.grad = tensor(-18.)\n", "x.grad = tensor(-12.)\n", "b.grad = tensor(-6.)\n", "y.grad = tensor(6.)\n" ] } ], "source": [ "L.backward()\n", "\n", "print(\"w.grad =\", w.grad)\n", "print(\"x.grad =\", x.grad)\n", "print(\"b.grad =\", b.grad)\n", "print(\"y.grad =\", y.grad)" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-06", "metadata": {}, "source": [ "That is PyTorch autograd. The forward pass created a graph; `.backward()` sent a seed gradient of $1$\n", "from $L$ through that graph in reverse.\n", "\n", "For comparison with the paper calculation, here is every stored value:" ] }, { "cell_type": "code", "execution_count": 3, "id": "scalar-autograd-simple-07", "metadata": { "purpose": "Print and verify the complete PyTorch value-and-gradient ledger" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "node value grad\n", "w 2.0 -18.0\n", "x 3.0 -12.0\n", "m 6.0 -6.0\n", "b 1.0 -6.0\n", "a 7.0 -6.0\n", "y 10.0 6.0\n", "e -3.0 -6.0\n", "L 9.0 1.0\n" ] } ], "source": [ "torch_nodes = {\"w\": w, \"x\": x, \"m\": m, \"b\": b,\n", " \"a\": a, \"y\": y, \"e\": e, \"L\": L}\n", "torch_reference = {\n", " name: (node.item(), node.grad.item())\n", " for name, node in torch_nodes.items()\n", "}\n", "\n", "print(f\"{'node':<5} {'value':>8} {'grad':>8}\")\n", "for name, (value, grad) in torch_reference.items():\n", " print(f\"{name:<5} {value:8.1f} {grad:8.1f}\")\n", "\n", "expected = {\n", " \"w\": (2, -18), \"x\": (3, -12), \"m\": (6, -6), \"b\": (1, -6),\n", " \"a\": (7, -6), \"y\": (10, 6), \"e\": (-3, -6), \"L\": (9, 1),\n", "}\n", "assert torch_reference == expected" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-08", "metadata": {}, "source": [ "## 2 · A tiny autograd engine from scratch\n", "\n", "Suppose one operation $f$ takes a value $u$ and produces $v=f(u)$. In this notebook's convention,\n", "the **forward construction** makes $u$ a direct **parent** (or operand) of $v$, while $v$ is the output\n", "(or child) created from $u$. These words describe one operation in the computation graph—not a whole\n", "neural-network layer. Libraries sometimes choose different names; here, follow `v.parents`, whose links\n", "point back to the direct operands.\n", "\n", "
\n", " \"Parent\n", "
\n", "

\n", " v.parents owns the ParentLink, and that link points back to u.\n", " The engine does not need u to keep a list of its children. On a phone, scroll sideways.\n", "

\n", "
\n", " Concrete link from our graph: in m = w × x, choose v = m and the parent u = w.\n", " Backward reads upstream m.grad = −6, reads local ∂m/∂w = x = 3 from that link,\n", " computes the temporary contribution −6 × 3 = −18, and adds it to w.grad.\n", " The other parent link uses u = x, local derivative w = 2, and contributes −12 to x.grad.\n", "
\n", "\n", "For the autograd calculation, a `Value` needs only:\n", "\n", "- its number in `data`,\n", "- its accumulated loss-gradient buffer in `grad`, and\n", "- ordered links to the operands that directly produced it.\n", "\n", "Our teaching class also stores `label` and `op` so diagrams can say “m” and “×”. They are display\n", "metadata: changing those strings does not change the forward number or any gradient.\n", "\n", "Each `ParentLink` has exactly two fields: `.value` points to the parent operand, and `.local_grad` holds\n", "the evaluated local derivative along that edge. For example, $m=wx$ remembers\n", "$(w,\\partial m/\\partial w=x)$ and $(x,\\partial m/\\partial x=w)$.\n", "\n", "The gradient names are always relative to the operation currently running:\n", "\n", "- upstream:\n", " $g_v=\\partial L/\\partial v$, already accumulated in `v.grad`;\n", "- local:\n", " $\\partial v/\\partial u$, stored in the link from output $v$ to parent $u$;\n", "- edge contribution to the parent\n", " (the “downstream contribution” in our color legend):\n", " $\\Delta g_u=g_v(\\partial v/\\partial u)$, computed during backward and added to `u.grad`.\n", "\n", "The edge contribution is temporary: only the accumulated result in `u.grad` remains. `w` is a leaf\n", "because it has no dependencies. `L` is the forward output or sink; `backward(L)` treats it as the\n", "starting node—the root of the reverse traversal.\n", "\n", "### Before writing `backward`: decide when a node is ready\n", "\n", "A node must not send its gradient to its parents until **all gradient contributions arriving at that\n", "node have been added to its `.grad` buffer**. We therefore need a dependency-safe processing order.\n", "Start at `L`, follow its saved `ParentLink`s toward the inputs, and append each node only after all its\n", "parents have been appended. Reversing the resulting list gives the safe order for backward.\n", "\n", "
\n", " \"Dependency-safe\n", "
\n", "

\n", " The formal name for any ordering that puts every dependency before the value that uses it is a\n", " topological order. We first understand the readiness rule; the name is secondary.\n", " On a phone, scroll sideways.\n", "

\n", "\n", "First apply the append-after-parents rule to one small part of the graph:\n", "\n", "```text\n", "visit(m):\n", " visit(w) → w has no parents → append w\n", " visit(x) → x has no parents → append x\n", " both parents are ready → append m\n", "```\n", "\n", "Starting from `L` applies that same rule recursively to the whole graph. With our stored parent order,\n", "the traversal produces exactly\n", "\n", "```text\n", "dependency-first: w, x, m, b, a, y, e, L\n", "process back: L, e, y, a, b, m, x, w\n", "```\n", "\n", "Other dependency-safe lists are possible—for example, two independent leaves can swap places.\n", "`seen` matters when a value feeds several later operations: following links from `L` may reach that same\n", "object more than once, but it must be appended and processed only once. `seen` deduplicates `Value`\n", "objects in the node schedule—not edges. If an operation is `w * w`, it still stores two `ParentLink`s,\n", "and backward still processes both contributions. For $r=w^2$ at $w=2$, `w` appears once in the node\n", "schedule, but the two links each contribute $2$; `w.grad += contribution` therefore gives $4$." ] }, { "cell_type": "code", "execution_count": 4, "id": "scalar-autograd-simple-09", "metadata": { "purpose": "Define parent links, local rules, a dependency-safe ordering helper, and backward" }, "outputs": [], "source": [ "class ParentLink:\n", " def __init__(self, value, local_grad):\n", " self.value = value\n", " self.local_grad = float(local_grad)\n", "\n", "class Value:\n", " def __init__(self, data, label=\"\", parents=(), op=\"\"):\n", " self.data = float(data)\n", " self.grad = 0.0\n", " self.label = label\n", " self.parents = tuple(parents)\n", " self.op = op\n", "\n", "def multiply(u, v, label):\n", " return Value(u.data * v.data, label,\n", " parents=(ParentLink(u, v.data),\n", " ParentLink(v, u.data)), op=\"×\")\n", "\n", "def add(u, v, label):\n", " return Value(u.data + v.data, label,\n", " parents=(ParentLink(u, 1.0),\n", " ParentLink(v, 1.0)), op=\"+\")\n", "\n", "def subtract(u, v, label):\n", " return Value(u.data - v.data, label,\n", " parents=(ParentLink(u, 1.0),\n", " ParentLink(v, -1.0)), op=\"−\")\n", "\n", "def square(u, label):\n", " return Value(u.data ** 2, label,\n", " parents=(ParentLink(u, 2 * u.data),), op=\"²\")\n", "\n", "def dependency_safe_order(root):\n", " \"\"\"Return reachable Values with every parent before its output.\"\"\"\n", " safe_order = []\n", " seen = set()\n", "\n", " def append_after_parents(node):\n", " # A shared Value may be reachable from the loss by several paths.\n", " # Visit and append that object only once.\n", " if id(node) in seen:\n", " return\n", " seen.add(id(node))\n", "\n", " # Follow output -> ParentLink -> operand, starting from the loss.\n", " for link in node.parents:\n", " append_after_parents(link.value)\n", "\n", " # Only now are all direct parents earlier in safe_order.\n", " safe_order.append(node)\n", "\n", " append_after_parents(root)\n", " return safe_order\n", "\n", "def backward(root):\n", " safe_order = dependency_safe_order(root)\n", "\n", " # 1. Clear old accumulated gradients, then seed the loss.\n", " for node in safe_order:\n", " node.grad = 0.0\n", " root.grad = 1.0\n", "\n", " # 2. Reverse the safe order. A node's full upstream gradient is\n", " # ready before that node sends contributions to its parents.\n", " steps = []\n", " for output in reversed(safe_order):\n", " # One saved parent link gives one chain-rule update.\n", " for link in output.parents:\n", " parent = link.value\n", " upstream = output.grad\n", " local = link.local_grad\n", " contribution = upstream * local\n", " before = parent.grad\n", " parent.grad += contribution\n", "\n", " # Keep a teaching trace; autograd only needs the update above.\n", " steps.append({\n", " \"output\": output.label,\n", " \"upstream\": upstream,\n", " \"parent\": parent.label,\n", " \"local\": local,\n", " \"downstream\": contribution,\n", " \"before\": before,\n", " \"after\": parent.grad,\n", " })\n", " return steps" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-10", "metadata": {}, "source": [ "The visible code separates **finding a safe order** from **doing the calculus**:\n", "\n", "1. `dependency_safe_order(root)` starts at `L`. `append_after_parents` follows each stored link to a\n", " direct operand and calls itself there first. Only after those calls return does it append the current\n", " node. `seen` makes a shared object a no-op on its second visit.\n", "2. `backward(root)` clears every reachable `.grad`, then seeds `L.grad = 1` because\n", " $\\partial L/\\partial L=1$.\n", "3. `reversed(safe_order)` processes `L, e, y, a, b, m, x, w`. At every saved link from output $v$\n", " to parent $u$, it reads the now-complete upstream gradient from `v.grad`, multiplies by the saved\n", " local derivative, and accumulates the result in `u.grad`.\n", "\n", "Leaves such as `w` still appear in the processing list. They simply have no parent links, so there is\n", "nothing further to update when their turn arrives.\n", "\n", "| quantity | where it lives |\n", "|---|---|\n", "| upstream $g_v$ | already accumulated in `v.grad` |\n", "| local $\\partial v/\\partial u$ | saved in `link.local_grad` during the forward pass |\n", "| edge contribution to parent $\\Delta g_u$ | temporary variable `contribution` for this one edge |\n", "| accumulated $g_u$ | updated in `parent.grad` |\n", "\n", "The autograd graph does **not** store a separate downstream gradient forever. It computes one edge\n", "contribution, adds it to the parent's buffer, and that buffer later becomes the upstream gradient for\n", "the parent. Our returned `steps` list is only a teaching log: it copies each contribution and the\n", "before/after values so we can display them.\n", "\n", "Three deliberate boundaries keep this engine small:\n", "\n", "- it assumes an acyclic computation graph (a DAG);\n", "- it seeds a scalar loss with `1`; vector outputs would need an explicit upstream seed;\n", "- it clears reachable `.grad` buffers at the start of every call. PyTorch normally **accumulates**\n", " gradients across `.backward()` calls until you clear them." ] }, { "cell_type": "code", "execution_count": 5, "id": "scalar-autograd-simple-11", "metadata": { "cellView": "form", "jupyter": { "source_hidden": true }, "purpose": "Define the compact graph, backward trace cards, and complete stored-state view" }, "outputs": [], "source": [ "#| echo: false\n", "import importlib.util\n", "import json\n", "import subprocess\n", "import sys\n", "import warnings\n", "from html import escape as escape_html\n", "\n", "if importlib.util.find_spec(\"graphviz\") is None:\n", " subprocess.run(\n", " [sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"graphviz\"],\n", " check=True,\n", " )\n", "\n", "from graphviz import Digraph\n", "from IPython.display import HTML, display\n", "\n", "def draw_graph(root, show_grad=True, min_width=780):\n", " nodes = dependency_safe_order(root)\n", "\n", " dot = Digraph(format=\"svg\")\n", " dot.attr(rankdir=\"LR\", bgcolor=\"transparent\", pad=\"0.15\",\n", " nodesep=\"0.25\", ranksep=\"0.45\")\n", " dot.attr(\"node\", fontname=\"Helvetica\", fontsize=\"11\")\n", " dot.attr(\"edge\", fontname=\"Helvetica\", fontsize=\"9\", color=\"#60777b\")\n", "\n", " for node in nodes:\n", " node_id = \"value_\" + node.label\n", " grad = f\"{node.grad:g}\" if show_grad else \"—\"\n", " fill = \"#e8f7f5\" if show_grad and node.grad != 0 else \"#ffffff\"\n", " border = \"#2C7A7B\" if show_grad and node.grad != 0 else \"#1f3a40\"\n", " dot.node(\n", " node_id,\n", " label=f\"{{ {node.label} | value {node.data:g} | grad {grad} }}\",\n", " shape=\"record\", style=\"rounded,filled\", fillcolor=fill,\n", " color=border, fontcolor=\"#1f3a40\",\n", " )\n", " if node.op:\n", " op_id = \"op_\" + node.label\n", " dot.node(op_id, label=node.op, shape=\"circle\",\n", " width=\"0.32\", height=\"0.32\", margin=\"0.03\",\n", " style=\"filled\", fillcolor=\"#2B6CB0\",\n", " color=\"#2B6CB0\", fontcolor=\"white\")\n", " dot.edge(op_id, node_id)\n", " for link in node.parents:\n", " local_label = (\n", " f\"∂{node.label}/∂{link.value.label}=\"\n", " f\"{link.local_grad:g}\"\n", " )\n", " dot.edge(\n", " \"value_\" + link.value.label,\n", " op_id,\n", " label=local_label,\n", " fontcolor=\"#2B6CB0\",\n", " )\n", " svg = dot.pipe(format=\"svg\").decode(\"utf-8\")\n", " svg = svg.replace(\n", " \"'\n", " f'
' + svg + '
'\n", " )\n", "\n", "def trace_dependency_safe_order(root):\n", " \"\"\"Run the real append-after-parents recursion while recording display states.\"\"\"\n", " safe_order = []\n", " seen_ids = set()\n", " seen_values = []\n", " call_stack = []\n", " events = []\n", "\n", " def snapshot(\n", " kind, *, node=None, parent=None, link_index=None,\n", " local_grad=None, line, message,\n", " ):\n", " event = {\n", " \"kind\": kind,\n", " \"stack\": [value.label for value in call_stack],\n", " \"seen\": [value.label for value in seen_values],\n", " \"order\": [value.label for value in safe_order],\n", " \"line\": line,\n", " \"message\": message,\n", " }\n", " if node is not None:\n", " event[\"node\"] = node.label\n", " if parent is not None:\n", " event[\"parent\"] = parent.label\n", " if link_index is not None:\n", " event[\"link_index\"] = link_index\n", " if local_grad is not None:\n", " event[\"local_grad\"] = float(local_grad)\n", " events.append(event)\n", "\n", " snapshot(\n", " \"ready\", line=8,\n", " message=(\n", " \"Press Next or Play to call \"\n", " \"append_after_parents(L).\"\n", " ),\n", " )\n", "\n", " def append_after_parents(node):\n", " call_stack.append(node)\n", " label = escape_html(node.label)\n", "\n", " if id(node) in seen_ids:\n", " snapshot(\n", " \"skip\", node=node, line=2,\n", " message=(\n", " f\"{label} is already in seen, so this repeated \"\n", " \"call returns without appending it again.\"\n", " ),\n", " )\n", " call_stack.pop()\n", " return\n", "\n", " seen_ids.add(id(node))\n", " seen_values.append(node)\n", " parent_count = len(node.parents)\n", " if parent_count:\n", " remaining = \"its parent\" if parent_count == 1 else f\"all {parent_count} parents\"\n", " enter_message = (\n", " f\"Enter {label}. It is new, so mark it seen; \"\n", " f\"this call must now visit {remaining}.\"\n", " )\n", " else:\n", " enter_message = (\n", " f\"Enter leaf {label}. It is new, so mark it seen. \"\n", " \"It has no ParentLinks, so it can be appended next.\"\n", " )\n", " snapshot(\"enter\", node=node, line=3, message=enter_message)\n", "\n", " for link_index, link in enumerate(node.parents):\n", " parent = link.value\n", " parent_label = escape_html(parent.label)\n", " snapshot(\n", " \"follow\", node=node, parent=parent,\n", " link_index=link_index, local_grad=link.local_grad, line=5,\n", " message=(\n", " f\"Follow {label}.parents[{link_index}] to \"\n", " f\"{parent_label}, then call \"\n", " f\"append_after_parents({parent_label}).\"\n", " ),\n", " )\n", " append_after_parents(parent)\n", " if link_index + 1 < parent_count:\n", " next_action = \"Continue to the next ParentLink.\"\n", " else:\n", " next_action = \"All of this node's parent calls are now finished.\"\n", " snapshot(\n", " \"unwind\", node=node, parent=parent, line=5,\n", " message=(\n", " f\"The call for {parent_label} has returned to \"\n", " f\"{label}. {next_action}\"\n", " ),\n", " )\n", "\n", " safe_order.append(node)\n", " snapshot(\n", " \"append\", node=node, line=6,\n", " message=(\n", " f\"Append {label} to safe_order. Every direct \"\n", " f\"parent of {label} is already earlier in the list.\"\n", " ),\n", " )\n", " call_stack.pop()\n", "\n", " append_after_parents(root)\n", " snapshot(\n", " \"return\", line=9,\n", " message=(\n", " \"The recursive helper is finished. Return the dependency-first \"\n", " \"safe_order unchanged.\"\n", " ),\n", " )\n", " snapshot(\n", " \"reverse\", line=11,\n", " message=(\n", " \"Now move into backward. Its loop reads \"\n", " \"reversed(safe_order) before performing any local-gradient arithmetic.\"\n", " ),\n", " )\n", "\n", " # This trace must be a faithful observation of the actual helper above it.\n", " assert safe_order == dependency_safe_order(root)\n", " assert events[-1][\"order\"] == [node.label for node in safe_order]\n", " return events\n", "\n", "def show_topological_sort_animation(root, events):\n", " \"\"\"Send the validated traversal to the isolated hosted interactive.\"\"\"\n", " nodes = dependency_safe_order(root)\n", " order_labels = [node.label for node in nodes]\n", " backward_labels = list(reversed(order_labels))\n", " parent_labels = {\n", " node.label: [link.value.label for link in node.parents]\n", " for node in nodes\n", " }\n", " node_data = {node.label: node.data for node in nodes}\n", " node_ops = {node.label: node.op for node in nodes}\n", " link_state = {\n", " node.label: [\n", " (link.value.label, float(link.local_grad))\n", " for link in node.parents\n", " ]\n", " for node in nodes\n", " }\n", "\n", " # The layout below is purpose-built for this lecture's exact graph.\n", " expected_order = [\"w\", \"x\", \"m\", \"b\", \"a\", \"y\", \"e\", \"L\"]\n", " expected_backward = [\"L\", \"e\", \"y\", \"a\", \"b\", \"m\", \"x\", \"w\"]\n", " expected_parents = {\n", " \"w\": [], \"x\": [], \"m\": [\"w\", \"x\"], \"b\": [],\n", " \"a\": [\"m\", \"b\"], \"y\": [], \"e\": [\"a\", \"y\"], \"L\": [\"e\"],\n", " }\n", " expected_data = {\n", " \"w\": 2.0, \"x\": 3.0, \"m\": 6.0, \"b\": 1.0,\n", " \"a\": 7.0, \"y\": 10.0, \"e\": -3.0, \"L\": 9.0,\n", " }\n", " expected_ops = {\n", " \"w\": \"\", \"x\": \"\", \"m\": \"×\", \"b\": \"\",\n", " \"a\": \"+\", \"y\": \"\", \"e\": \"−\", \"L\": \"²\",\n", " }\n", " expected_links = {\n", " \"w\": [], \"x\": [], \"m\": [(\"w\", 3.0), (\"x\", 2.0)], \"b\": [],\n", " \"a\": [(\"m\", 1.0), (\"b\", 1.0)], \"y\": [],\n", " \"e\": [(\"a\", 1.0), (\"y\", -1.0)], \"L\": [(\"e\", -6.0)],\n", " }\n", " assert order_labels == expected_order\n", " assert backward_labels == expected_backward\n", " assert parent_labels == expected_parents\n", " assert node_data == expected_data\n", " assert node_ops == expected_ops\n", " assert link_state == expected_links\n", " assert events[0][\"kind\"] == \"ready\"\n", " assert events[-2][\"kind\"] == \"return\"\n", " assert events[-1][\"kind\"] == \"reverse\"\n", " assert events[-1][\"order\"] == expected_order\n", " assert sum(event[\"kind\"] == \"enter\" for event in events) == 8\n", " assert sum(event[\"kind\"] == \"follow\" for event in events) == 7\n", " assert sum(event[\"kind\"] == \"unwind\" for event in events) == 7\n", " assert sum(event[\"kind\"] == \"append\" for event in events) == 8\n", " assert sum(event[\"kind\"] == \"return\" for event in events) == 1\n", " assert sum(event[\"kind\"] == \"reverse\" for event in events) == 1\n", " assert not any(event[\"kind\"] == \"skip\" for event in events)\n", "\n", " # Audit every recorded intermediate state, not only the final list.\n", " previous_seen = []\n", " previous_order = []\n", " for event in events:\n", " seen = event[\"seen\"]\n", " order = event[\"order\"]\n", " stack = event[\"stack\"]\n", " assert seen[:len(previous_seen)] == previous_seen\n", " assert order[:len(previous_order)] == previous_order\n", " assert len(seen) == len(set(seen))\n", " assert len(order) == len(set(order))\n", " assert set(order) <= set(seen)\n", " assert set(stack) <= set(seen)\n", " for output_label, parent_label in zip(stack, stack[1:]):\n", " assert parent_label in parent_labels[output_label]\n", "\n", " if event[\"kind\"] == \"enter\":\n", " assert stack[-1] == event[\"node\"] == seen[-1]\n", " elif event[\"kind\"] == \"follow\":\n", " assert stack[-1] == event[\"node\"]\n", " index = event[\"link_index\"]\n", " assert 0 <= index < len(link_state[event[\"node\"]])\n", " expected_parent, expected_local = link_state[event[\"node\"]][index]\n", " assert event[\"parent\"] == expected_parent\n", " assert event[\"local_grad\"] == expected_local\n", " elif event[\"kind\"] == \"append\":\n", " assert stack[-1] == event[\"node\"] == order[-1]\n", " assert set(parent_labels[event[\"node\"]]) <= set(order[:-1])\n", " elif event[\"kind\"] == \"unwind\":\n", " assert stack[-1] == event[\"node\"]\n", " assert event[\"parent\"] in parent_labels[event[\"node\"]]\n", " elif event[\"kind\"] in {\"return\", \"reverse\"}:\n", " assert not stack and order == expected_order\n", "\n", " previous_seen = seen\n", " previous_order = order\n", "\n", " payload = {\n", " \"events\": events,\n", " \"parents\": parent_labels,\n", " \"order\": order_labels,\n", " \"backward\": backward_labels,\n", " }\n", " hosted_url = (\n", " \"https://nipunbatra.github.io/interactive-articles/\"\n", " \"autograd-topological-order/\"\n", " )\n", " hosted_embed_url = hosted_url + \"?embed=1\"\n", " payload_json = json.dumps(\n", " payload, ensure_ascii=False, separators=(\",\", \":\")\n", " ).replace(\"\"\n", " \".scalar-topology-embed{margin:16px 0 22px;border:1px solid #D8E1E2;\"\n", " \"border-radius:16px;overflow:hidden;background:#fff;color:#1F3A40;\"\n", " \"font-family:Manrope,ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif}\"\n", " \".scalar-topology-embed__bar{display:flex;align-items:center;justify-content:space-between;\"\n", " \"gap:16px;padding:12px 15px;border-bottom:1px solid #E2E9EA;background:#F7FAFA}\"\n", " \".scalar-topology-embed__title{font-weight:800;line-height:1.25}\"\n", " \".scalar-topology-embed__hint{display:block;margin-top:2px;color:#52696D;\"\n", " \"font-size:.84rem;font-weight:500}\"\n", " \".scalar-topology-embed__open{flex:0 0 auto;padding:8px 11px;border:1px solid #2C7A7B;\"\n", " \"border-radius:8px;color:#1F6666!important;background:#fff;text-decoration:none!important;\"\n", " \"font-size:.86rem;font-weight:750}\"\n", " \".scalar-topology-embed__open:hover{background:#E8F7F5}\"\n", " \".scalar-topology-animation-frame{display:block;width:100%;height:1080px;border:0;background:#fff}\"\n", " \".scalar-topology-embed__fallback{margin:0;padding:9px 15px;border-top:1px solid #E2E9EA;\"\n", " \"color:#52696D;font-size:.84rem}\"\n", " \".scalar-topology-embed__print{display:none}\"\n", " \"@media(max-width:850px){.scalar-topology-animation-frame{height:1780px}}\"\n", " \"@media(max-width:640px){\"\n", " \".scalar-topology-embed__bar{align-items:flex-start;flex-direction:column;gap:9px}\"\n", " \".scalar-topology-animation-frame{height:2400px}\"\n", " \"}\"\n", " \"@media print{\"\n", " \".scalar-topology-embed__bar,.scalar-topology-animation-frame,\"\n", " \".scalar-topology-embed__fallback{display:none!important}\"\n", " \".scalar-topology-embed__print{display:block;padding:14px 16px;line-height:1.5}\"\n", " \"}\"\n", " \"\"\n", " '
'\n", " '
'\n", " '
Dependency-safe ordering, step by step
'\n", " 'Uses the live graph and ParentLinks built above.'\n", " '
'\n", " f'Open full screen
'\n", " ''\n", " '

If the embedded view does not load, '\n", " f'open the interactive '\n", " 'in a new tab.

'\n", " '
Dependency-safe order: '\n", " 'w → x → m → b → a → y → e → L
'\n", " 'Backward schedule: L → e → y → a → b → m → x → w
'\n", " '
'\n", " \"\"\n", " )\n", " with warnings.catch_warnings():\n", " warnings.filterwarnings(\n", " \"ignore\",\n", " message=\"Consider using IPython.display.IFrame instead\",\n", " category=UserWarning,\n", " )\n", " display(HTML(iframe))\n", "\n", "def show_backward_steps(steps, highlight_outputs=()):\n", " rows = []\n", " for number, step in enumerate(steps, start=1):\n", " output = step[\"output\"]\n", " parent = step[\"parent\"]\n", " focused = output in highlight_outputs\n", " border = \"#2B6CB0\" if focused else \"#d6e2e1\"\n", " background = \"#f4f8ff\" if focused else \"#fff\"\n", " rows.append(\n", " f\"
\"\n", " f\"
{number}. {output} → {parent}
\"\n", " \"
\"\n", " f\"g{output} = {step['upstream']:g}\"\n", " \"  ×  \"\n", " f\"∂{output}/∂{parent} = {step['local']:g}\"\n", " \"  =  \"\n", " f\"Δg{parent} = {step['downstream']:g}\"\n", " \"
\"\n", " f\"
{parent}.grad: \"\n", " f\"{step['before']:g} → {step['after']:g}
\"\n", " )\n", "\n", " display(HTML(\n", " \"
\"\n", " \"Seed: L.grad = ∂L/∂L = 1
\" + \"\".join(rows)\n", " ))\n", "\n", "def show_complete_state(root):\n", " \"\"\"Show every stored Value field and every saved ParentLink.\"\"\"\n", " cards = []\n", " for node in dependency_safe_order(root):\n", " label = escape_html(node.label)\n", " label_value = escape_html(repr(node.label))\n", " op_value = escape_html(repr(node.op))\n", " op_note = \"\" if node.op else \"  —  input / leaf\"\n", "\n", " if node.parents:\n", " parent_rows = []\n", " for index, link in enumerate(node.parents, start=1):\n", " parent = escape_html(link.value.label)\n", " parent_rows.append(\n", " \"
\"\n", " f\".parents[{index - 1}]\"\n", " f\"ParentLink(value={parent})\"\n", " \".value\"\n", " f\"{parent}\"\n", " \".local_grad\"\n", " f\"{link.local_grad:g}\"\n", " f\"   (= ∂{label}/∂{parent})
\"\n", " )\n", " parents_html = \"\".join(parent_rows)\n", " else:\n", " parents_html = (\n", " \"
\"\n", " \".parents = ()  —  no direct operands
\"\n", " )\n", "\n", " cards.append(\n", " \"
\"\n", " \"
\"\n", " f\"Value {label}\"\n", " f\".label = {label_value}
\"\n", " \"
\"\n", " \"
\"\n", " \".data\"\n", " f\"{node.data:g}\"\n", " \".grad\"\n", " f\"{node.grad:g} \"\n", " f\"(= ∂L/∂{label})\"\n", " \".op\"\n", " f\"{op_value}{op_note}\"\n", " \".parents\"\n", " f\"{len(node.parents)} saved link\"\n", " f\"{'s' if len(node.parents) != 1 else ''}
\"\n", " f\"
{parents_html}
\"\n", " )\n", "\n", " order_text = \" → \".join(escape_html(node.label) for node in dependency_safe_order(root))\n", " display(HTML(\n", " \"
\"\n", " \"
\"\n", " \"Complete stored state after backward
\"\n", " \"The graph above stays compact. These cards expose every \"\n", " \"Value field and every saved ParentLink. \"\n", " \"The orange edge contribution is not a Value or ParentLink field; \"\n", " \"it survives only in the optional steps teaching trace.
\"\n", " f\"Displayed in dependency-first order: {order_text}\"\n", " \"
\"\n", " \"
\" + \"\".join(cards) + \"
\"\n", " ))" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-12", "metadata": {}, "source": [ "Build the **same forward graph**, one readable line per operation. The interactive trace immediately\n", "below is generated from these actual `Value` objects and their stored `ParentLink`s—not from a separate\n", "hand-written event list. On a phone, scroll the graph sideways:" ] }, { "cell_type": "code", "execution_count": 6, "id": "scalar-autograd-simple-13", "metadata": { "purpose": "Inspect saved links, draw the forward graph, and step through the actual ordering traversal" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "What m=wx stored during forward:\n", " parent w: local ∂m/∂w = 3\n", " parent x: local ∂m/∂x = 2\n", "\n", "Dependency-first order: w → x → m → b → a → y → e → L\n", "Backward will process: L → e → y → a → b → m → x → w\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "value_w\n", "\n", "w\n", "\n", "value 2\n", "\n", "grad —\n", "\n", "\n", "\n", "op_m\n", "\n", "×\n", "\n", "\n", "\n", "value_w->op_m\n", "\n", "\n", "∂m/∂w=3\n", "\n", "\n", "\n", "value_x\n", "\n", "x\n", "\n", "value 3\n", "\n", "grad —\n", "\n", "\n", "\n", "value_x->op_m\n", "\n", "\n", "∂m/∂x=2\n", "\n", "\n", "\n", "value_m\n", "\n", "m\n", "\n", "value 6\n", "\n", "grad —\n", "\n", "\n", "\n", "op_a\n", "\n", "+\n", "\n", "\n", "\n", "value_m->op_a\n", "\n", "\n", "∂a/∂m=1\n", "\n", "\n", "\n", "op_m->value_m\n", "\n", "\n", "\n", "\n", "\n", "value_b\n", "\n", "b\n", "\n", "value 1\n", "\n", "grad —\n", "\n", "\n", "\n", "value_b->op_a\n", "\n", "\n", "∂a/∂b=1\n", "\n", "\n", "\n", "value_a\n", "\n", "a\n", "\n", "value 7\n", "\n", "grad —\n", "\n", "\n", "\n", "op_e\n", "\n", "−\n", "\n", "\n", "\n", "value_a->op_e\n", "\n", "\n", "∂e/∂a=1\n", "\n", "\n", "\n", "op_a->value_a\n", "\n", "\n", "\n", "\n", "\n", "value_y\n", "\n", "y\n", "\n", "value 10\n", "\n", "grad —\n", "\n", "\n", "\n", "value_y->op_e\n", "\n", "\n", "∂e/∂y=-1\n", "\n", "\n", "\n", "value_e\n", "\n", "e\n", "\n", "value -3\n", "\n", "grad —\n", "\n", "\n", "\n", "op_L\n", "\n", "²\n", "\n", "\n", "\n", "value_e->op_L\n", "\n", "\n", "∂L/∂e=-6\n", "\n", "\n", "\n", "op_e->value_e\n", "\n", "\n", "\n", "\n", "\n", "value_L\n", "\n", "L\n", "\n", "value 9\n", "\n", "grad —\n", "\n", "\n", "\n", "op_L->value_L\n", "\n", "\n", "\n", "\n", "\n", "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
Dependency-safe ordering, step by step
Uses the live graph and ParentLinks built above.
Open full screen

If the embedded view does not load, open the interactive in a new tab.

Dependency-safe order: w → x → m → b → a → y → e → L
Backward schedule: L → e → y → a → b → m → x → w
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "sw = Value(2.0, label=\"w\")\n", "sx = Value(3.0, label=\"x\")\n", "sb = Value(1.0, label=\"b\")\n", "sy = Value(10.0, label=\"y\")\n", "\n", "sm = multiply(sw, sx, \"m\")\n", "sa = add(sm, sb, \"a\")\n", "se = subtract(sa, sy, \"e\")\n", "sL = square(se, \"L\")\n", "\n", "print(\"What m=wx stored during forward:\")\n", "for link in sm.parents:\n", " print(f\" parent {link.value.label}: local ∂m/∂{link.value.label} = {link.local_grad:g}\")\n", "\n", "safe_order = dependency_safe_order(sL)\n", "print(\"\\nDependency-first order:\", \" → \".join(node.label for node in safe_order))\n", "print(\"Backward will process: \", \" → \".join(node.label for node in reversed(safe_order)))\n", "\n", "display(draw_graph(sL, show_grad=False))\n", "\n", "topology_events = trace_dependency_safe_order(sL)\n", "show_topological_sort_animation(sL, topology_events)" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-14", "metadata": {}, "source": [ "Now run backward once, then inspect the result at three levels:\n", "\n", "1. the **edge-by-edge trace** shows every chain-rule multiplication and accumulation;\n", "2. the **compact graph** shows the whole computation without overcrowding it;\n", "3. the **complete state cards** expose every field on every `Value`, including every saved parent link.\n", "\n", "Only `steps = backward(sL)` performs differentiation. The two `show_...` helpers and `draw_graph` are\n", "teaching displays; removing them would not change any gradient." ] }, { "cell_type": "code", "execution_count": 7, "id": "scalar-autograd-simple-15", "metadata": { "purpose": "Run backward, show every edge, redraw gradients, and expose complete stored state" }, "outputs": [ { "data": { "text/html": [ "
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -6  =  Δge = -6
e.grad: 0 → -6
2. e → a
ge = -6  ×  ∂e/∂a = 1  =  Δga = -6
a.grad: 0 → -6
3. e → y
ge = -6  ×  ∂e/∂y = -1  =  Δgy = 6
y.grad: 0 → 6
4. a → m
ga = -6  ×  ∂a/∂m = 1  =  Δgm = -6
m.grad: 0 → -6
5. a → b
ga = -6  ×  ∂a/∂b = 1  =  Δgb = -6
b.grad: 0 → -6
6. m → w
gm = -6  ×  ∂m/∂w = 3  =  Δgw = -18
w.grad: 0 → -18
7. m → x
gm = -6  ×  ∂m/∂x = 2  =  Δgx = -12
x.grad: 0 → -12
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "value_w\n", "\n", "w\n", "\n", "value 2\n", "\n", "grad -18\n", "\n", "\n", "\n", "op_m\n", "\n", "×\n", "\n", "\n", "\n", "value_w->op_m\n", "\n", "\n", "∂m/∂w=3\n", "\n", "\n", "\n", "value_x\n", "\n", "x\n", "\n", "value 3\n", "\n", "grad -12\n", "\n", "\n", "\n", "value_x->op_m\n", "\n", "\n", "∂m/∂x=2\n", "\n", "\n", "\n", "value_m\n", "\n", "m\n", "\n", "value 6\n", "\n", "grad -6\n", "\n", "\n", "\n", "op_a\n", "\n", "+\n", "\n", "\n", "\n", "value_m->op_a\n", "\n", "\n", "∂a/∂m=1\n", "\n", "\n", "\n", "op_m->value_m\n", "\n", "\n", "\n", "\n", "\n", "value_b\n", "\n", "b\n", "\n", "value 1\n", "\n", "grad -6\n", "\n", "\n", "\n", "value_b->op_a\n", "\n", "\n", "∂a/∂b=1\n", "\n", "\n", "\n", "value_a\n", "\n", "a\n", "\n", "value 7\n", "\n", "grad -6\n", "\n", "\n", "\n", "op_e\n", "\n", "−\n", "\n", "\n", "\n", "value_a->op_e\n", "\n", "\n", "∂e/∂a=1\n", "\n", "\n", "\n", "op_a->value_a\n", "\n", "\n", "\n", "\n", "\n", "value_y\n", "\n", "y\n", "\n", "value 10\n", "\n", "grad 6\n", "\n", "\n", "\n", "value_y->op_e\n", "\n", "\n", "∂e/∂y=-1\n", "\n", "\n", "\n", "value_e\n", "\n", "e\n", "\n", "value -3\n", "\n", "grad -6\n", "\n", "\n", "\n", "op_L\n", "\n", "²\n", "\n", "\n", "\n", "value_e->op_L\n", "\n", "\n", "∂L/∂e=-6\n", "\n", "\n", "\n", "op_e->value_e\n", "\n", "\n", "\n", "\n", "\n", "value_L\n", "\n", "L\n", "\n", "value 9\n", "\n", "grad 1\n", "\n", "\n", "\n", "op_L->value_L\n", "\n", "\n", "\n", "\n", "\n", "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
Complete stored state after backward
The graph above stays compact. These cards expose every Value field and every saved ParentLink. The orange edge contribution is not a Value or ParentLink field; it survives only in the optional steps teaching trace.
Displayed in dependency-first order: w → x → m → b → a → y → e → L
Value w.label = 'w'
.data2.grad-18 (= ∂L/∂w).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value x.label = 'x'
.data3.grad-12 (= ∂L/∂x).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value m.label = 'm'
.data6.grad-6 (= ∂L/∂m).op'×'.parents2 saved links
.parents[0]ParentLink(value=w).valuew.local_grad3   (= ∂m/∂w)
.parents[1]ParentLink(value=x).valuex.local_grad2   (= ∂m/∂x)
Value b.label = 'b'
.data1.grad-6 (= ∂L/∂b).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value a.label = 'a'
.data7.grad-6 (= ∂L/∂a).op'+'.parents2 saved links
.parents[0]ParentLink(value=m).valuem.local_grad1   (= ∂a/∂m)
.parents[1]ParentLink(value=b).valueb.local_grad1   (= ∂a/∂b)
Value y.label = 'y'
.data10.grad6 (= ∂L/∂y).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value e.label = 'e'
.data-3.grad-6 (= ∂L/∂e).op'−'.parents2 saved links
.parents[0]ParentLink(value=a).valuea.local_grad1   (= ∂e/∂a)
.parents[1]ParentLink(value=y).valuey.local_grad-1   (= ∂e/∂y)
Value L.label = 'L'
.data9.grad1 (= ∂L/∂L).op'²'.parents1 saved link
.parents[0]ParentLink(value=e).valuee.local_grad-6   (= ∂L/∂e)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "steps = backward(sL)\n", "show_backward_steps(steps)\n", "\n", "display(draw_graph(sL, show_grad=True))\n", "show_complete_state(sL)" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-16", "metadata": {}, "source": [ "The trace contains every reverse edge. For example, the square sends $-6$ into `e.grad`. On the next\n", "operation, that same stored number becomes the upstream gradient $g_e$ for subtraction.\n", "\n", "A single row's product is one **edge contribution to `parent.grad`**—the quantity colored orange in our\n", "legend. If several paths return to one value, each row adds into the same buffer; only their sum is the\n", "full gradient at that parent.\n", "\n", "Finally, check that our tiny engine and PyTorch agree at every named value." ] }, { "cell_type": "code", "execution_count": 8, "id": "scalar-autograd-simple-17", "metadata": { "purpose": "Verify exact scratch-to-PyTorch parity" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "✓ Every value and gradient matches PyTorch.\n" ] } ], "source": [ "scratch_nodes = {\"w\": sw, \"x\": sx, \"m\": sm, \"b\": sb,\n", " \"a\": sa, \"y\": sy, \"e\": se, \"L\": sL}\n", "\n", "for name, node in scratch_nodes.items():\n", " torch_value, torch_grad = torch_reference[name]\n", " assert node.data == torch_value\n", " assert node.grad == torch_grad\n", "\n", "print(\"✓ Every value and gradient matches PyTorch.\")" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-18", "metadata": {}, "source": [ "## 3 · One neuron: fused sigmoid or atomic sigmoid?\n", "\n", "Now use a slightly larger graph:\n", "\n", "$$\n", "m=wx,\\qquad z=m+b,\\qquad s=\\sigma(z),\\qquad e=s-y,\\qquad L=e^2.\n", "$$\n", "\n", "Choose $w=0.5$, $x=2$, $b=-1$, and $y=1$. Then $z=0$, $s=0.5$, and $L=0.25$,\n", "so the backward numbers stay readable.\n", "\n", "We will build the sigmoid in two ways:\n", "\n", "- **fused autograd primitive:** one operation $s=\\sigma(z)$;\n", "- **atomic graph:** $n=-z$, $q=\\exp(n)$, $d=1+q$, and $s=1/d$.\n", "\n", "“Fused” here describes the autograd graph: several local steps are packaged behind one operation node.\n", "It does not mean that we are skipping the chain rule." ] }, { "cell_type": "code", "execution_count": 9, "id": "scalar-autograd-simple-19", "metadata": { "purpose": "Define the four atomic sigmoid rules and one fused sigmoid rule" }, "outputs": [], "source": [ "import math\n", "\n", "def negate(u, label):\n", " return Value(-u.data, label,\n", " parents=(ParentLink(u, -1.0),), op=\"−\")\n", "\n", "def exponential(u, label):\n", " out = math.exp(u.data)\n", " return Value(out, label,\n", " parents=(ParentLink(u, out),), op=\"exp\")\n", "\n", "def plus_one(u, label):\n", " return Value(1.0 + u.data, label,\n", " parents=(ParentLink(u, 1.0),), op=\"+1\")\n", "\n", "def reciprocal(u, label):\n", " return Value(1.0 / u.data, label,\n", " parents=(ParentLink(u, -1.0 / u.data**2),), op=\"1/x\")\n", "\n", "def sigmoid(u, label):\n", " # Stable forward formula; backward reuses the saved output s.\n", " if u.data >= 0:\n", " s = 1.0 / (1.0 + math.exp(-u.data))\n", " else:\n", " exp_z = math.exp(u.data)\n", " s = exp_z / (1.0 + exp_z)\n", " return Value(s, label,\n", " parents=(ParentLink(u, s * (1.0 - s)),), op=\"σ\")" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-20", "metadata": {}, "source": [ "The fused rule stores one local derivative:\n", "\n", "$$\n", "\\frac{\\partial s}{\\partial z}=s(1-s).\n", "$$\n", "\n", "The atomic graph stores four local derivatives. Their product is the same quantity:\n", "\n", "$$\n", "\\underbrace{\\left(-\\frac{1}{d^2}\\right)}_{s=1/d}\n", "\\underbrace{(1)}_{d=1+q}\n", "\\underbrace{(q)}_{q=\\exp(n)}\n", "\\underbrace{(-1)}_{n=-z}\n", "=\\frac{q}{d^2}=s(1-s).\n", "$$" ] }, { "cell_type": "code", "execution_count": 10, "id": "scalar-autograd-simple-21", "metadata": { "purpose": "Build both sigmoid graphs and show every backward edge" }, "outputs": [ { "data": { "text/html": [ "

Fused sigmoid · 8 reverse edges

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "value_w\n", "\n", "w\n", "\n", "value 0.5\n", "\n", "grad -0.5\n", "\n", "\n", "\n", "op_m\n", "\n", "×\n", "\n", "\n", "\n", "value_w->op_m\n", "\n", "\n", "∂m/∂w=2\n", "\n", "\n", "\n", "value_x\n", "\n", "x\n", "\n", "value 2\n", "\n", "grad -0.125\n", "\n", "\n", "\n", "value_x->op_m\n", "\n", "\n", "∂m/∂x=0.5\n", "\n", "\n", "\n", "value_m\n", "\n", "m\n", "\n", "value 1\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "op_z\n", "\n", "+\n", "\n", "\n", "\n", "value_m->op_z\n", "\n", "\n", "∂z/∂m=1\n", "\n", "\n", "\n", "op_m->value_m\n", "\n", "\n", "\n", "\n", "\n", "value_b\n", "\n", "b\n", "\n", "value -1\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "value_b->op_z\n", "\n", "\n", "∂z/∂b=1\n", "\n", "\n", "\n", "value_z\n", "\n", "z\n", "\n", "value 0\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "op_s\n", "\n", "σ\n", "\n", "\n", "\n", "value_z->op_s\n", "\n", "\n", "∂s/∂z=0.25\n", "\n", "\n", "\n", "op_z->value_z\n", "\n", "\n", "\n", "\n", "\n", "value_s\n", "\n", "s\n", "\n", "value 0.5\n", "\n", "grad -1\n", "\n", "\n", "\n", "op_e\n", "\n", "−\n", "\n", "\n", "\n", "value_s->op_e\n", "\n", "\n", "∂e/∂s=1\n", "\n", "\n", "\n", "op_s->value_s\n", "\n", "\n", "\n", "\n", "\n", "value_y\n", "\n", "y\n", "\n", "value 1\n", "\n", "grad 1\n", "\n", "\n", "\n", "value_y->op_e\n", "\n", "\n", "∂e/∂y=-1\n", "\n", "\n", "\n", "value_e\n", "\n", "e\n", "\n", "value -0.5\n", "\n", "grad -1\n", "\n", "\n", "\n", "op_L\n", "\n", "²\n", "\n", "\n", "\n", "value_e->op_L\n", "\n", "\n", "∂L/∂e=-1\n", "\n", "\n", "\n", "op_e->value_e\n", "\n", "\n", "\n", "\n", "\n", "value_L\n", "\n", "L\n", "\n", "value 0.25\n", "\n", "grad 1\n", "\n", "\n", "\n", "op_L->value_L\n", "\n", "\n", "\n", "\n", "\n", "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -1  =  Δge = -1
e.grad: 0 → -1
2. e → s
ge = -1  ×  ∂e/∂s = 1  =  Δgs = -1
s.grad: 0 → -1
3. e → y
ge = -1  ×  ∂e/∂y = -1  =  Δgy = 1
y.grad: 0 → 1
4. s → z
gs = -1  ×  ∂s/∂z = 0.25  =  Δgz = -0.25
z.grad: 0 → -0.25
5. z → m
gz = -0.25  ×  ∂z/∂m = 1  =  Δgm = -0.25
m.grad: 0 → -0.25
6. z → b
gz = -0.25  ×  ∂z/∂b = 1  =  Δgb = -0.25
b.grad: 0 → -0.25
7. m → w
gm = -0.25  ×  ∂m/∂w = 2  =  Δgw = -0.5
w.grad: 0 → -0.5
8. m → x
gm = -0.25  ×  ∂m/∂x = 0.5  =  Δgx = -0.125
x.grad: 0 → -0.125
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "

Atomic sigmoid · 11 reverse edges

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "value_w\n", "\n", "w\n", "\n", "value 0.5\n", "\n", "grad -0.5\n", "\n", "\n", "\n", "op_m\n", "\n", "×\n", "\n", "\n", "\n", "value_w->op_m\n", "\n", "\n", "∂m/∂w=2\n", "\n", "\n", "\n", "value_x\n", "\n", "x\n", "\n", "value 2\n", "\n", "grad -0.125\n", "\n", "\n", "\n", "value_x->op_m\n", "\n", "\n", "∂m/∂x=0.5\n", "\n", "\n", "\n", "value_m\n", "\n", "m\n", "\n", "value 1\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "op_z\n", "\n", "+\n", "\n", "\n", "\n", "value_m->op_z\n", "\n", "\n", "∂z/∂m=1\n", "\n", "\n", "\n", "op_m->value_m\n", "\n", "\n", "\n", "\n", "\n", "value_b\n", "\n", "b\n", "\n", "value -1\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "value_b->op_z\n", "\n", "\n", "∂z/∂b=1\n", "\n", "\n", "\n", "value_z\n", "\n", "z\n", "\n", "value 0\n", "\n", "grad -0.25\n", "\n", "\n", "\n", "op_n\n", "\n", "−\n", "\n", "\n", "\n", "value_z->op_n\n", "\n", "\n", "∂n/∂z=-1\n", "\n", "\n", "\n", "op_z->value_z\n", "\n", "\n", "\n", "\n", "\n", "value_n\n", "\n", "n\n", "\n", "value -0\n", "\n", "grad 0.25\n", "\n", "\n", "\n", "op_q\n", "\n", "exp\n", "\n", "\n", "\n", "value_n->op_q\n", "\n", "\n", "∂q/∂n=1\n", "\n", "\n", "\n", "op_n->value_n\n", "\n", "\n", "\n", "\n", "\n", "value_q\n", "\n", "q\n", "\n", "value 1\n", "\n", "grad 0.25\n", "\n", "\n", "\n", "op_d\n", "\n", "+1\n", "\n", "\n", "\n", "value_q->op_d\n", "\n", "\n", "∂d/∂q=1\n", "\n", "\n", "\n", "op_q->value_q\n", "\n", "\n", "\n", "\n", "\n", "value_d\n", "\n", "d\n", "\n", "value 2\n", "\n", "grad 0.25\n", "\n", "\n", "\n", "op_s\n", "\n", "1/x\n", "\n", "\n", "\n", "value_d->op_s\n", "\n", "\n", "∂s/∂d=-0.25\n", "\n", "\n", "\n", "op_d->value_d\n", "\n", "\n", "\n", "\n", "\n", "value_s\n", "\n", "s\n", "\n", "value 0.5\n", "\n", "grad -1\n", "\n", "\n", "\n", "op_e\n", "\n", "−\n", "\n", "\n", "\n", "value_s->op_e\n", "\n", "\n", "∂e/∂s=1\n", "\n", "\n", "\n", "op_s->value_s\n", "\n", "\n", "\n", "\n", "\n", "value_y\n", "\n", "y\n", "\n", "value 1\n", "\n", "grad 1\n", "\n", "\n", "\n", "value_y->op_e\n", "\n", "\n", "∂e/∂y=-1\n", "\n", "\n", "\n", "value_e\n", "\n", "e\n", "\n", "value -0.5\n", "\n", "grad -1\n", "\n", "\n", "\n", "op_L\n", "\n", "²\n", "\n", "\n", "\n", "value_e->op_L\n", "\n", "\n", "∂L/∂e=-1\n", "\n", "\n", "\n", "op_e->value_e\n", "\n", "\n", "\n", "\n", "\n", "value_L\n", "\n", "L\n", "\n", "value 0.25\n", "\n", "grad 1\n", "\n", "\n", "\n", "op_L->value_L\n", "\n", "\n", "\n", "\n", "\n", "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -1  =  Δge = -1
e.grad: 0 → -1
2. e → s
ge = -1  ×  ∂e/∂s = 1  =  Δgs = -1
s.grad: 0 → -1
3. e → y
ge = -1  ×  ∂e/∂y = -1  =  Δgy = 1
y.grad: 0 → 1
4. s → d
gs = -1  ×  ∂s/∂d = -0.25  =  Δgd = 0.25
d.grad: 0 → 0.25
5. d → q
gd = 0.25  ×  ∂d/∂q = 1  =  Δgq = 0.25
q.grad: 0 → 0.25
6. q → n
gq = 0.25  ×  ∂q/∂n = 1  =  Δgn = 0.25
n.grad: 0 → 0.25
7. n → z
gn = 0.25  ×  ∂n/∂z = -1  =  Δgz = -0.25
z.grad: 0 → -0.25
8. z → m
gz = -0.25  ×  ∂z/∂m = 1  =  Δgm = -0.25
m.grad: 0 → -0.25
9. z → b
gz = -0.25  ×  ∂z/∂b = 1  =  Δgb = -0.25
b.grad: 0 → -0.25
10. m → w
gm = -0.25  ×  ∂m/∂w = 2  =  Δgw = -0.5
w.grad: 0 → -0.5
11. m → x
gm = -0.25  ×  ∂m/∂x = 0.5  =  Δgx = -0.125
x.grad: 0 → -0.125
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "def build_sigmoid_neuron(*, fused):\n", " nodes = {\n", " \"w\": Value(0.5, label=\"w\"),\n", " \"x\": Value(2.0, label=\"x\"),\n", " \"b\": Value(-1.0, label=\"b\"),\n", " \"y\": Value(1.0, label=\"y\"),\n", " }\n", " nodes[\"m\"] = multiply(nodes[\"w\"], nodes[\"x\"], \"m\")\n", " nodes[\"z\"] = add(nodes[\"m\"], nodes[\"b\"], \"z\")\n", "\n", " if fused:\n", " nodes[\"s\"] = sigmoid(nodes[\"z\"], \"s\")\n", " else:\n", " nodes[\"n\"] = negate(nodes[\"z\"], \"n\")\n", " nodes[\"q\"] = exponential(nodes[\"n\"], \"q\")\n", " nodes[\"d\"] = plus_one(nodes[\"q\"], \"d\")\n", " nodes[\"s\"] = reciprocal(nodes[\"d\"], \"s\")\n", "\n", " nodes[\"e\"] = subtract(nodes[\"s\"], nodes[\"y\"], \"e\")\n", " nodes[\"L\"] = square(nodes[\"e\"], \"L\")\n", " return nodes\n", "\n", "fused_nodes = build_sigmoid_neuron(fused=True)\n", "atomic_nodes = build_sigmoid_neuron(fused=False)\n", "fused_steps = backward(fused_nodes[\"L\"])\n", "atomic_steps = backward(atomic_nodes[\"L\"])\n", "\n", "display(HTML(\"

Fused sigmoid · 8 reverse edges

\"))\n", "display(draw_graph(fused_nodes[\"L\"], show_grad=True, min_width=1180))\n", "show_backward_steps(fused_steps, highlight_outputs={\"s\"})\n", "\n", "display(HTML(\"

Atomic sigmoid · 11 reverse edges

\"))\n", "display(draw_graph(atomic_nodes[\"L\"], show_grad=True, min_width=1700))\n", "show_backward_steps(atomic_steps, highlight_outputs={\"s\", \"d\", \"q\", \"n\"})" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-22", "metadata": {}, "source": [ "The blue-highlighted cards are the only part that changed:\n", "\n", "- fused sigmoid: one update, $g_z=g_s\\,s(1-s)=(-1)(0.25)=-0.25$;\n", "- atomic sigmoid: four updates, ending with the same $g_z=-0.25$.\n", "\n", "Fusion therefore gives a smaller graph and fewer intermediate gradient buffers. A real library can also\n", "use a numerically stable sigmoid implementation. The mathematics is unchanged: the single fused local\n", "derivative is exactly the product of the four atomic local derivatives." ] }, { "cell_type": "code", "execution_count": 11, "id": "scalar-autograd-simple-23", "metadata": { "purpose": "Verify fused-to-atomic equivalence and PyTorch parity" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "✓ Fused, atomic, and PyTorch agree.\n", " sigmoid local: 4 atomic factors = 1 fused factor = 0.25\n", " final gradients: w = -0.5, x = -0.125, b = -0.25, y = 1\n" ] } ], "source": [ "common = (\"w\", \"x\", \"m\", \"b\", \"z\", \"s\", \"y\", \"e\", \"L\")\n", "for name in common:\n", " assert math.isclose(fused_nodes[name].data, atomic_nodes[name].data)\n", " assert math.isclose(fused_nodes[name].grad, atomic_nodes[name].grad)\n", "\n", "assert [(s[\"output\"], s[\"parent\"]) for s in fused_steps] == [\n", " (\"L\", \"e\"), (\"e\", \"s\"), (\"e\", \"y\"), (\"s\", \"z\"),\n", " (\"z\", \"m\"), (\"z\", \"b\"), (\"m\", \"w\"), (\"m\", \"x\"),\n", "]\n", "assert [(s[\"output\"], s[\"parent\"]) for s in atomic_steps] == [\n", " (\"L\", \"e\"), (\"e\", \"s\"), (\"e\", \"y\"), (\"s\", \"d\"),\n", " (\"d\", \"q\"), (\"q\", \"n\"), (\"n\", \"z\"), (\"z\", \"m\"),\n", " (\"z\", \"b\"), (\"m\", \"w\"), (\"m\", \"x\"),\n", "]\n", "\n", "expected_sigmoid = {\n", " \"w\": (0.5, -0.5), \"x\": (2.0, -0.125), \"m\": (1.0, -0.25),\n", " \"b\": (-1.0, -0.25), \"z\": (0.0, -0.25), \"s\": (0.5, -1.0),\n", " \"y\": (1.0, 1.0), \"e\": (-0.5, -1.0), \"L\": (0.25, 1.0),\n", "}\n", "for name, (value, grad) in expected_sigmoid.items():\n", " assert math.isclose(fused_nodes[name].data, value)\n", " assert math.isclose(fused_nodes[name].grad, grad)\n", "\n", "fused_local = next(\n", " step[\"local\"] for step in fused_steps\n", " if step[\"output\"] == \"s\" and step[\"parent\"] == \"z\"\n", ")\n", "atomic_locals = [\n", " step[\"local\"] for step in atomic_steps\n", " if step[\"output\"] in {\"s\", \"d\", \"q\", \"n\"}\n", "]\n", "assert math.isclose(math.prod(atomic_locals), fused_local)\n", "\n", "tw = torch.tensor(0.5, requires_grad=True)\n", "tx = torch.tensor(2.0, requires_grad=True)\n", "tb = torch.tensor(-1.0, requires_grad=True)\n", "ty = torch.tensor(1.0, requires_grad=True)\n", "tL = (torch.sigmoid(tw * tx + tb) - ty) ** 2\n", "tL.backward()\n", "\n", "assert math.isclose(fused_nodes[\"w\"].grad, tw.grad.item())\n", "assert math.isclose(fused_nodes[\"x\"].grad, tx.grad.item())\n", "assert math.isclose(fused_nodes[\"b\"].grad, tb.grad.item())\n", "assert math.isclose(fused_nodes[\"y\"].grad, ty.grad.item())\n", "\n", "print(\"✓ Fused, atomic, and PyTorch agree.\")\n", "print(\" sigmoid local: 4 atomic factors = 1 fused factor =\", fused_local)\n", "print(\" final gradients: w = -0.5, x = -0.125, b = -0.25, y = 1\")" ] }, { "cell_type": "markdown", "id": "scalar-autograd-simple-24", "metadata": {}, "source": [ "## Takeaway\n", "\n", "Both systems do the same three things:\n", "\n", "1. run the forward operations and store parent links plus local derivatives,\n", "2. start with $g_L=1$ in the loss's `.grad` buffer,\n", "3. compute upstream\n", " $\\times$ local\n", " $=$ edge contribution to the parent, then add it to\n", " the parent's `.grad` buffer.\n", "\n", "Our tiny `Value` record and local rules make those steps visible. Fusion does not change the calculus;\n", "it packages a product of local derivatives behind one operation. PyTorch generalizes these ideas to\n", "tensors, neural-network layers, accelerators, and large models." ] } ], "metadata": { "course": { "evidence": "exact scalar example; traced backward; fused-versus-atomic sigmoid parity", "lecture": 4, "title": "Computation Graphs, Backpropagation & Autograd" }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.2" } }, "nbformat": 4, "nbformat_minor": 5 }