{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-1 fw=all -->\n",
    "\n",
    "# Matrix Calculus and Automatic Differentiation\n",
    "\n",
    "The previous two sections built differentiation up from a single weight\n",
    "(that section) to the gradient of a scalar loss over\n",
    "many weights (that section). Real network layers,\n",
    "however, map *vectors to vectors*, and their parameters are *matrices*, so the\n",
    "natural object of study is the derivative of a vector-valued map, the *Jacobian*,\n",
    "and the natural question is why the backward pass — PyTorch's `loss.backward()`,\n",
    "JAX's `grad`, and their kin — is cheap. This section answers\n",
    "both: **backpropagation is reverse-mode automatic differentiation, a sequence of\n",
    "vector--Jacobian products**, and the choice between forward- and reverse-mode AD\n",
    "is dictated by the *shape* of the Jacobian you are after. Along the way we\n",
    "re-derive the handful of matrix identities that actually recur in deep learning,\n",
    "and we build both flavours of automatic differentiation from scratch in a few\n",
    "dozen lines of Python, so that the framework's autograd stops being magic.\n",
    "\n",
    "We use the following imports throughout this section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-imports"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "import numpy as np\n",
    "import jax\n",
    "from jax import numpy as jnp"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-2 fw=all -->\n",
    "\n",
    "## Derivatives of Vector- and Matrix-Valued Maps\n",
    "\n",
    "A layer $\\mathbf f:\\mathbb R^n\\to\\mathbb R^m$ does not have \"a derivative\" but a\n",
    "whole matrix of them. The single object that organizes them, and generalizes both\n",
    "the slope of that section and the gradient of\n",
    "that section in one stroke, is the *Jacobian*.\n",
    "\n",
    "### The Jacobian as the Best Linear Approximation\n",
    "\n",
    "In one variable, the derivative is the slope of the best straight-line fit:\n",
    "$f(x+\\delta)\\approx f(x)+f'(x)\\,\\delta$. The vector-valued generalization keeps\n",
    "this idea verbatim. We say $\\mathbf f$ is *differentiable* at $\\mathbf x$ if there\n",
    "is a matrix $\\mathbf J\\in\\mathbb R^{m\\times n}$ such that\n",
    "\n",
    "$$\n",
    "\\mathbf f(\\mathbf x+\\boldsymbol\\delta)\n",
    "   = \\mathbf f(\\mathbf x) + \\mathbf J\\,\\boldsymbol\\delta + o(\\|\\boldsymbol\\delta\\|),\n",
    "$$\n",
    "\n",
    "i.e. the error of the linear model $\\mathbf f(\\mathbf x)+\\mathbf J\\boldsymbol\\delta$\n",
    "shrinks *faster* than $\\boldsymbol\\delta$ itself as $\\boldsymbol\\delta\\to\\mathbf 0$.\n",
    "That matrix, if it exists, is unique, is called the *Jacobian*\n",
    "$\\mathbf J_{\\mathbf f}(\\mathbf x)$, and its entries are exactly the partial\n",
    "derivatives:\n",
    "\n",
    "$$\n",
    "[\\mathbf J_{\\mathbf f}]_{ij} = \\frac{\\partial f_i}{\\partial x_j}.\n",
    "$$\n",
    "\n",
    "To see why, feed the equation the perturbation\n",
    "$\\boldsymbol\\delta=t\\,\\mathbf e_j$ (a small step along axis $j$) and read off\n",
    "component $i$: dividing by $t$ and letting $t\\to0$ leaves precisely\n",
    "$\\partial f_i/\\partial x_j$ in the $(i,j)$ slot. So row $i$ of $\\mathbf J$ collects\n",
    "the partials of the $i$-th output, and column $j$ records how *all* outputs respond\n",
    "to nudging input $j$. The Jacobian *is* the best local linear approximation; the\n",
    "partial-derivative formula is a consequence, not the definition.\n",
    "\n",
    "the figure makes the definition visible. Up close, a\n",
    "differentiable map *is* a linear map: a small circle of inputs around\n",
    "$\\mathbf x_0$ lands on (very nearly) an ellipse: the image of that circle under\n",
    "$\\mathbf J(\\mathbf x_0)$, exactly the picture from\n",
    "that section of what a matrix does to the\n",
    "plane. The leftover bend is the $o(\\|\\boldsymbol\\delta\\|)$ remainder, which\n",
    "vanishes faster than the circle shrinks.\n",
    "\n",
    "![Up close, a differentiable map is a linear map. The nonlinear $\\mathbf f(x,y)=(x+\\sin y,\\ y+x^2/2)$ carries a small circle and grid around $\\mathbf x_0=(0.5,0.5)$ (left) into the output plane (right). The true image (blue) is nearly indistinguishable from the ellipse $\\mathbf f(\\mathbf x_0)+\\mathbf J(\\mathbf x_0)\\,\\boldsymbol\\delta$ predicted by the Jacobian (orange); the small mismatch is the $o(\\|\\boldsymbol\\delta\\|)$ remainder and shrinks faster than the circle does.](../img/mdl-cal-jacobian-ellipse.svg)\n",
    "\n",
    "Two special cases recover everything we have already met. When $m=1$ (a scalar\n",
    "field $f:\\mathbb R^n\\to\\mathbb R$) the Jacobian is a *single row*, the *row*\n",
    "gradient $\\partial f/\\partial\\mathbf x = [\\partial f/\\partial x_1,\\ldots,\\partial\n",
    "f/\\partial x_n]$, and the equation is the first-order\n",
    "expansion $f(\\mathbf x+\\boldsymbol\\delta)\\approx f(\\mathbf x)+\\mathbf J\\boldsymbol\\delta$\n",
    "of that section. And the Jacobian of the *gradient field*\n",
    "$\\nabla f:\\mathbb R^n\\to\\mathbb R^n$ is the Hessian\n",
    "$\\mathbf H=\\mathbf J_{\\nabla f}$ of the equation: differentiating once\n",
    "more turns the row gradient into the matrix of second partials, where the entry\n",
    "$[\\mathbf J_{\\nabla f}]_{ij}=\\partial^2 f/\\partial x_j\\partial x_i$ matches the\n",
    "$(i,j)$ Hessian entry because mixed partials commute (Clairaut's theorem, proved\n",
    "in that section). One construction, three familiar\n",
    "objects.\n",
    "\n",
    "Let us pin this down on a concrete $\\mathbb R^2\\to\\mathbb R^2$ map and check the\n",
    "linear approximation against a finite difference, the numerical signature of a\n",
    "correct derivative."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-jacobian-finite-diff"
   },
   "source": [
    "def f(v):                                 # f: R^2 -> R^2\n",
    "    x, y = v\n",
    "    return np.array([x**2 * y, np.sin(x + y)])\n",
    "\n",
    "def J(v):                                 # the exact Jacobian, entry by entry\n",
    "    x, y = v\n",
    "    return np.array([[2 * x * y, x**2],\n",
    "                     [np.cos(x + y), np.cos(x + y)]])\n",
    "\n",
    "x0 = np.array([1.0, 0.5])\n",
    "delta = np.array([1e-3, -2e-3])\n",
    "linear = f(x0) + J(x0) @ delta            # first-order prediction\n",
    "exact = f(x0 + delta)                     # true value\n",
    "print('linear approx:', linear.round(6))\n",
    "print('exact        :', exact.round(6))\n",
    "print('error / |delta|:', np.linalg.norm(exact - linear) / np.linalg.norm(delta))"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "linear approx: [0.499    0.997424]\nexact        : [0.498996 0.997424]\nerror / |delta|: 0.0015819444529976167\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-3 fw=all -->\n",
    "\n",
    "The error is tiny *relative to* $\\|\\boldsymbol\\delta\\|$, and shrinking it by a\n",
    "further factor of ten in $\\boldsymbol\\delta$ shrinks the error by a factor of a\n",
    "hundred, the $o(\\|\\boldsymbol\\delta\\|)$ signature of the equation.\n",
    "\n",
    "### The Chain Rule Is Jacobian Composition\n",
    "\n",
    "The multivariate chain rule of that section, a sum\n",
    "over paths, is in matrix form the *multiplication of Jacobians*.\n",
    "\n",
    "**Proposition (chain rule).** *If $\\mathbf f:\\mathbb R^n\\to\\mathbb R^p$ is\n",
    "differentiable at $\\mathbf x$ and $\\mathbf g:\\mathbb R^p\\to\\mathbb R^m$ is\n",
    "differentiable at $\\mathbf f(\\mathbf x)$, then $\\mathbf g\\circ\\mathbf f$ is\n",
    "differentiable at $\\mathbf x$ and*\n",
    "\n",
    "$$\n",
    "\\mathbf J_{\\mathbf g\\circ\\mathbf f}(\\mathbf x)\n",
    "   = \\mathbf J_{\\mathbf g}\\!\\bigl(\\mathbf f(\\mathbf x)\\bigr)\\,\\mathbf J_{\\mathbf f}(\\mathbf x).\n",
    "$$\n",
    "\n",
    "**Proof.** Write $\\mathbf y=\\mathbf f(\\mathbf x)$ and let $\\mathbf A=\\mathbf\n",
    "J_{\\mathbf f}(\\mathbf x)$, $\\mathbf B=\\mathbf J_{\\mathbf g}(\\mathbf y)$. Linearity\n",
    "is the entire argument: substitute the linear model for $\\mathbf f$ into the linear\n",
    "model for $\\mathbf g$. By the equation,\n",
    "$\\mathbf f(\\mathbf x+\\boldsymbol\\delta)=\\mathbf y+\\mathbf A\\boldsymbol\\delta\n",
    "+o(\\|\\boldsymbol\\delta\\|)$. Calling that displacement\n",
    "$\\boldsymbol\\eta=\\mathbf A\\boldsymbol\\delta+o(\\|\\boldsymbol\\delta\\|)$ and feeding it\n",
    "to $\\mathbf g$,\n",
    "\n",
    "$$\n",
    "\\mathbf g\\bigl(\\mathbf f(\\mathbf x+\\boldsymbol\\delta)\\bigr)\n",
    "   = \\mathbf g(\\mathbf y) + \\mathbf B\\,\\boldsymbol\\eta + o(\\|\\boldsymbol\\eta\\|)\n",
    "   = \\mathbf g(\\mathbf y) + \\mathbf B\\mathbf A\\,\\boldsymbol\\delta + o(\\|\\boldsymbol\\delta\\|),\n",
    "$$\n",
    "\n",
    "because $\\mathbf B$ applied to the $o(\\|\\boldsymbol\\delta\\|)$ remainder is again\n",
    "$o(\\|\\boldsymbol\\delta\\|)$, and $\\|\\boldsymbol\\eta\\|=O(\\|\\boldsymbol\\delta\\|)$\n",
    "(bounded by a constant multiple of $\\|\\boldsymbol\\delta\\|$) folds\n",
    "$\\mathbf g$'s own remainder into $o(\\|\\boldsymbol\\delta\\|)$ as well. The\n",
    "coefficient of $\\boldsymbol\\delta$ is the best linear map, so by uniqueness it *is*\n",
    "the Jacobian of the composite: $\\mathbf J_{\\mathbf g\\circ\\mathbf f}=\\mathbf B\\mathbf A$.\n",
    "$\\blacksquare$\n",
    "\n",
    "The entry-wise chain rule falls out by reading off one entry of $\\mathbf B\\mathbf A$:\n",
    "$[\\mathbf B\\mathbf A]_{ik}=\\sum_j \\partial g_i/\\partial y_j\\cdot\\partial f_j/\\partial x_k$\n",
    "is exactly the \"sum over intermediate variables\" of\n",
    "that section. Matrix multiplication *is* the\n",
    "bookkeeping of paths.\n",
    "\n",
    "Iterating the equation is where deep learning enters. A network\n",
    "of depth $L$ is a composition\n",
    "$\\mathbf f=\\mathbf f_L\\circ\\cdots\\circ\\mathbf f_1$, so its end-to-end Jacobian is\n",
    "one long matrix product,\n",
    "\n",
    "$$\n",
    "\\mathbf J = \\mathbf J_L\\,\\mathbf J_{L-1}\\cdots\\mathbf J_1,\n",
    "$$\n",
    "\n",
    "each factor the local Jacobian of one layer. This is the same product-of-Jacobians\n",
    "that governs backpropagation through time in recurrent networks\n",
    "(that section).\n",
    "\n",
    "Two shapes dominate those factors. A dense layer $\\mathbf x\\mapsto\\mathbf W\\mathbf x$ is its own best\n",
    "linear approximation, so its Jacobian is just $\\mathbf W$. An *elementwise*\n",
    "activation $\\boldsymbol\\varphi(\\mathbf x)=(\\varphi(x_1),\\ldots,\\varphi(x_n))$\n",
    "couples no pair of coordinates (output $i$ depends only on input $i$), so every\n",
    "off-diagonal partial vanishes and its Jacobian is *diagonal*,\n",
    "$\\mathbf J_{\\boldsymbol\\varphi}(\\mathbf x)\n",
    "=\\operatorname{diag}\\bigl(\\varphi'(x_1),\\ldots,\\varphi'(x_n)\\bigr)$.\n",
    "Multiplying by it is mere elementwise scaling; for ReLU, $\\varphi'$ is a mask of\n",
    "zeros and ones marking the active units. A standard deep network alternates the\n",
    "two, so the equation specializes to\n",
    "$\\mathbf J=\\operatorname{diag}(\\boldsymbol\\varphi_L')\\,\\mathbf W_L\\cdots\n",
    "\\operatorname{diag}(\\boldsymbol\\varphi_1')\\,\\mathbf W_1$: dense weight matrices\n",
    "interleaved with cheap diagonal masks. That diagonal factor is the single most\n",
    "common Jacobian in a real network, and it is why an activation's backward pass\n",
    "costs only $O(n)$.\n",
    "\n",
    "Matrix multiplication is\n",
    "*associative*, so we may evaluate the equation in any order.\n",
    "Multiplying right-to-left propagates *inputs forward*; multiplying left-to-right\n",
    "propagates *sensitivities backward*. Those two parenthesizations are precisely\n",
    "forward-mode and reverse-mode automatic differentiation, and choosing between them\n",
    "is just choosing the cheaper way to multiply a chain of matrices. (They are the\n",
    "two *extreme* orderings; finding the cheapest ordering for a general computation\n",
    "graph (*optimal Jacobian accumulation*) is NP-complete [@Naumann.2008],\n",
    "so frameworks ship only the two ends; see [@Baydin.Pearlmutter.Radul.ea.2018]\n",
    "for a survey.)\n",
    "\n",
    "### Layout Conventions: Numerator vs Denominator\n",
    "\n",
    "One bookkeeping decision needs settling before we derive anything. For a scalar\n",
    "field $f:\\mathbb R^n\\to\\mathbb R$, do we write the derivative as a *row* (matching the Jacobian\n",
    "the equation) or as a *column* (matching the input vector\n",
    "$\\mathbf x$)? Both are in use, and they are *transposes of each other*:\n",
    "\n",
    "* **Numerator (Jacobian) layout.** The derivative has the shape of the\n",
    "  *numerator*: $\\partial\\mathbf y/\\partial\\mathbf x$ is $m\\times n$, so the scalar\n",
    "  case is a row. This makes the chain rule the equation read\n",
    "  left-to-right with no transposes, which is why it is the convention of choice for\n",
    "  the Jacobian theory above and for MIT's *Matrix Calculus for Machine Learning\n",
    "  and Beyond* course (18.S096, Edelman and Johnson).\n",
    "* **Denominator (gradient) layout.** The derivative has the shape of the\n",
    "  *denominator*: $\\partial f/\\partial\\mathbf x$ is a *column* matching $\\mathbf x$.\n",
    "  This is the statistician's $\\nabla_{\\mathbf x} f$, the thing you subtract in\n",
    "  gradient descent, and the convention of the derivations below.\n",
    "\n",
    "Switching layouts simply transposes the result, so the gradient column and the\n",
    "Jacobian row of a scalar field are related by\n",
    "\n",
    "$$\n",
    "\\nabla_{\\mathbf x} f = \\Bigl(\\frac{\\partial f}{\\partial\\mathbf x}\\Bigr)^{\\!\\top}.\n",
    "$$\n",
    "\n",
    "**This book's convention.** We use **numerator layout for Jacobians** of genuine\n",
    "vector-valued maps (the chain rule then composes without transposes) and\n",
    "**denominator layout for gradients**, writing $\\nabla_{\\mathbf x} f$ for the column you feed an optimizer.\n",
    "Whenever the two could be confused we make the shape explicit. With that settled,\n",
    "the identities in the next section are gradients, hence columns; if you ever land\n",
    "on an expression whose shape does not match its denominator, you have a missing\n",
    "transpose; use the equation to fix it. The Matrix\n",
    "Cookbook [@Petersen.Pedersen.ea.2008] tabulates hundreds of such identities;\n",
    "the point of the next section is that you need to *memorize* almost none of them.\n",
    "\n",
    "## A Few Key Identities, Derived Not Tabulated\n",
    "\n",
    "A reference table of matrix-derivative identities is long and forgettable. In\n",
    "practice only a handful recur in deep learning, and each one yields to the same\n",
    "two-step recipe of Parr and Howard's primer [@Parr.Howard.2018]:\n",
    "*differentiate one component with the ordinary scalar rules, then reassemble the\n",
    "components into a matrix expression.* A faster cross-check, due to the same source,\n",
    "is the **scalar-collapse heuristic**: a correct matrix identity must reduce to the\n",
    "familiar single-variable result when every matrix is $1\\times1$ (where products are\n",
    "scalar products, sums are sums, and transposes do nothing), so you can often *guess*\n",
    "the matrix form from the scalar one and fix the shapes by inserting transposes.\n",
    "We derive four identities this way; all are gradients, hence columns (denominator\n",
    "layout, the equation).\n",
    "\n",
    "**The linear form $\\nabla_{\\mathbf x}\\,\\mathbf a^\\top\\mathbf x = \\mathbf a$.** With\n",
    "$f(\\mathbf x)=\\mathbf a^\\top\\mathbf x=\\sum_i a_i x_i$, the $k$-th partial keeps only\n",
    "the $i=k$ term, so $\\partial f/\\partial x_k = a_k$; stacking these gives\n",
    "$\\nabla_{\\mathbf x} f=\\mathbf a$. The scalar shadow is $\\frac{d}{dx}(ax)=a$, exactly\n",
    "the matrix answer with the transpose doing nothing. The one thing to\n",
    "keep straight is *which* variable we differentiate by: it is $x_k$, not $a_k$.\n",
    "\n",
    "**The quadratic form $\\nabla_{\\mathbf x}\\,\\mathbf x^\\top\\mathbf A\\mathbf x =\n",
    "(\\mathbf A+\\mathbf A^\\top)\\mathbf x$.** Write the form in Einstein notation\n",
    "(that section) as\n",
    "$\\mathbf x^\\top\\mathbf A\\mathbf x = x_i a_{ij} x_j$. The product rule on\n",
    "$\\partial/\\partial x_k$ hits both $x$'s,\n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial x_k}\\bigl(x_i a_{ij} x_j\\bigr)\n",
    "   = \\delta_{ik}\\,a_{ij}x_j + x_i a_{ij}\\,\\delta_{jk}\n",
    "   = a_{kj}x_j + x_i a_{ik}\n",
    "   = (a_{ki}+a_{ik})\\,x_i,\n",
    "$$\n",
    "\n",
    "where the last step uses the fact that a summation (dummy) index can be renamed\n",
    "freely: the first term is a sum over $j$, the second a sum over $i$, so renaming\n",
    "$j\\to i$ in the first term ($a_{kj}x_j = a_{ki}x_i$) puts both sums over the same\n",
    "index $i$ and lets them combine. The bracket $a_{ki}+a_{ik}$ is the\n",
    "$(k,i)$ entry of $\\mathbf A+\\mathbf A^\\top$, so the $k$-th partial is the $k$-th\n",
    "entry of $(\\mathbf A+\\mathbf A^\\top)\\mathbf x$, giving\n",
    "$\\nabla_{\\mathbf x}\\,\\mathbf x^\\top\\mathbf A\\mathbf x=(\\mathbf A+\\mathbf A^\\top)\\mathbf x$.\n",
    "The scalar collapse is $\\frac{d}{dx}(a x^2)=2ax=(a+a)x$; the matrix answer is the\n",
    "same with $\\mathbf A^\\top$ standing in for the \"other copy\" of $a$. For *symmetric*\n",
    "$\\mathbf A$ this simplifies to the much-used $2\\mathbf A\\mathbf x$.\n",
    "\n",
    "Derive, then verify: both identities should survive the same autograd audit we\n",
    "will shortly apply to softmax, at a randomly drawn point."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-matrix-calculus-autodiff-a-few-key-identities-derived-not-tabulated-1"
   },
   "source": [
    "# Autograd audit of the two identities at a random point\n",
    "k1, k2, k3 = jax.random.split(jax.random.PRNGKey(0), 3)\n",
    "a, A = jax.random.normal(k1, (3,)), jax.random.normal(k2, (3, 3))\n",
    "x = jax.random.normal(k3, (3,))\n",
    "print('grad a^T x   equals a        :',\n",
    "      bool(jnp.allclose(jax.grad(lambda x: a @ x)(x), a)))\n",
    "print('grad x^T A x equals (A+A^T)x :',\n",
    "      bool(jnp.allclose(jax.grad(lambda x: x @ A @ x)(x), (A + A.T) @ x)))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "grad a^T x   equals a        : True\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "grad x^T A x equals (A+A^T)x : True\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-4 fw=all -->\n",
    "\n",
    "**The least-squares gradient $\\nabla_{\\mathbf W}\\|\\mathbf W\\mathbf x-\\mathbf y\\|^2$.**\n",
    "Let $\\mathbf r=\\mathbf W\\mathbf x-\\mathbf y$ be the residual, so the loss is\n",
    "$\\mathbf r^\\top\\mathbf r=\\sum_i r_i^2$ with $r_i=\\sum_j W_{ij}x_j-y_i$. Then\n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial W_{ab}}\\sum_i r_i^2\n",
    "   = 2\\sum_i r_i\\,\\frac{\\partial r_i}{\\partial W_{ab}}\n",
    "   = 2 r_a x_b = 2\\,[\\mathbf r\\mathbf x^\\top]_{ab},\n",
    "$$\n",
    "\n",
    "since $\\partial r_i/\\partial W_{ab}=\\delta_{ia}x_b$. Assembling the $(a,b)$ entries\n",
    "into a matrix the shape of $\\mathbf W$,\n",
    "\n",
    "$$\n",
    "\\nabla_{\\mathbf W}\\|\\mathbf W\\mathbf x-\\mathbf y\\|^2 = 2\\,(\\mathbf W\\mathbf x-\\mathbf y)\\,\\mathbf x^\\top,\n",
    "$$\n",
    "\n",
    "a rank-one *outer product* of the residual with the input; this is the gradient a\n",
    "single linear layer hands back during backpropagation. The companion gradient with\n",
    "respect to the *input*, $\\nabla_{\\mathbf x}\\|\\mathbf W\\mathbf x-\\mathbf y\\|^2\n",
    "=2\\mathbf W^\\top(\\mathbf W\\mathbf x-\\mathbf y)$, follows identically; note the\n",
    "$\\mathbf W^\\top$. The same pattern, a residual hit by the transpose of whatever\n",
    "multiplies the variable, persists when everything becomes a matrix. For the\n",
    "factorization loss $\\|\\mathbf X-\\mathbf U\\mathbf V\\|_F^2=\\sum_{ij}R_{ij}^2$ with\n",
    "residual $\\mathbf R=\\mathbf X-\\mathbf U\\mathbf V$, one entry of $\\mathbf V$\n",
    "appears in $R_{ij}$ through $\\partial R_{ij}/\\partial V_{ab}\n",
    "=-U_{ia}\\,\\delta_{jb}$, so\n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial V_{ab}}\\sum_{ij}R_{ij}^2\n",
    "   = -2\\sum_{i} U_{ia}R_{ib}\n",
    "   = -2\\,[\\mathbf U^\\top\\mathbf R]_{ab},\n",
    "\\qquad\\text{i.e.}\\qquad\n",
    "\\nabla_{\\mathbf V}\\|\\mathbf X-\\mathbf U\\mathbf V\\|_F^2\n",
    "   = -2\\,\\mathbf U^\\top(\\mathbf X-\\mathbf U\\mathbf V),\n",
    "$$\n",
    "\n",
    "exactly the $\\mathbf W^\\top$-shaped answer the scalar-collapse heuristic predicts\n",
    "from $\\frac{d}{dv}(x-uv)^2=-2u(x-uv)$.\n",
    "\n",
    "**The softmax--cross-entropy Jacobian.** The derivative a classifier evaluates\n",
    "most often also has the simplest form. Let $\\mathbf p=\\operatorname{softmax}(\\mathbf z)$,\n",
    "$p_i=e^{z_i}/\\sum_k e^{z_k}$. Quotient rule on $\\partial p_i/\\partial z_j$ gives the\n",
    "two cases $p_i(1-p_i)$ when $i=j$ and $-p_i p_j$ when $i\\neq j$, which combine into\n",
    "\n",
    "$$\n",
    "\\frac{\\partial p_i}{\\partial z_j} = p_i(\\delta_{ij}-p_j),\n",
    "\\qquad\n",
    "\\mathbf J_{\\operatorname{softmax}} = \\operatorname{diag}(\\mathbf p) - \\mathbf p\\mathbf p^\\top.\n",
    "$$\n",
    "\n",
    "Now compose with the cross-entropy loss $\\ell=-\\sum_i y_i\\log p_i$ against a\n",
    "one-hot target $\\mathbf y$. The chain rule the equation multiplies\n",
    "the loss's row gradient $\\partial\\ell/\\partial\\mathbf p$ by\n",
    "the equation, and the algebra collapses:\n",
    "\n",
    "$$\n",
    "\\frac{\\partial\\ell}{\\partial z_j}\n",
    "   = -\\sum_i \\frac{y_i}{p_i}\\,p_i(\\delta_{ij}-p_j)\n",
    "   = -y_j + p_j\\sum_i y_i\n",
    "   = p_j - y_j,\n",
    "$$\n",
    "\n",
    "using $\\sum_i y_i=1$. The gradient of softmax--cross-entropy with respect to the\n",
    "*logits* (the raw scores $\\mathbf z$ that softmax turns into probabilities) is\n",
    "just $\\mathbf p-\\mathbf y$: predicted probability minus truth. The softmax\n",
    "Jacobian cancels against the cross-entropy gradient. Libraries also fuse the two\n",
    "into a single primitive for a numerical reason: computing $\\log p_i$ from an\n",
    "already-evaluated softmax takes the logarithm of a value that may have underflowed\n",
    "to zero, whereas the fused primitive computes\n",
    "$\\log p_i = z_i - \\log\\sum_k e^{z_k}$ directly through a stable log-sum-exp. Let\n",
    "us confirm both the equation and the cancellation against a\n",
    "framework's autograd."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-softmax-jacobian"
   },
   "source": [
    "z = jnp.array([1.0, -0.5, 2.0])\n",
    "p = jax.nn.softmax(z)\n",
    "J = jnp.diag(p) - jnp.outer(p, p)              # our formula\n",
    "J_ad = jax.jacobian(jax.nn.softmax)(z)\n",
    "y = jnp.array([0.0, 1.0, 0.0])                 # one-hot target (class 1)\n",
    "loss = lambda z: -(y * jnp.log(jax.nn.softmax(z))).sum()\n",
    "print('softmax Jacobian matches autograd:', bool(jnp.allclose(J, J_ad, atol=1e-6)))\n",
    "print('d loss / d z :', np.asarray(jax.grad(loss)(z)).round(6))\n",
    "print('p - y        :', np.asarray(p - y).round(6))"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "softmax Jacobian matches autograd: True\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "d loss / d z : [ 0.253716 -0.943388  0.689672]\np - y        : [ 0.253716 -0.943388  0.689672]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-5 fw=all -->\n",
    "\n",
    "The logit gradient lands exactly on $\\mathbf p-\\mathbf y$, no matter how the\n",
    "intermediate Jacobian looks. Four identities, one method: differentiate a single\n",
    "component, then read the matrix back off the indices, sanity-checking against the\n",
    "scalar collapse.\n",
    "\n",
    "## Forward-Mode AD and Dual Numbers\n",
    "\n",
    "We now stop differentiating by hand. Automatic differentiation evaluates a function\n",
    "*and* its derivative in one sweep, to full numerical precision, by carrying\n",
    "derivative information alongside every value. It is neither symbolic differentiation\n",
    "(which manipulates formulas and explodes in size, as the swamp of repeated terms in\n",
    "that section showed) nor numerical differentiation\n",
    "(finite differences, which trade off truncation against round-off error). It is\n",
    "the chain rule, applied mechanically to the *program* that computes the\n",
    "function [@Baydin.Pearlmutter.Radul.ea.2018]. *Forward mode* is the simplest\n",
    "incarnation, and its algebra is that of *dual numbers*.\n",
    "\n",
    "### Dual Numbers: an Algebra That Carries Derivatives\n",
    "\n",
    "Define a *dual number* $a+b\\,\\varepsilon$ by analogy with the complex numbers, but\n",
    "with a new symbol $\\varepsilon$ satisfying\n",
    "\n",
    "$$\n",
    "\\varepsilon^2 = 0, \\qquad \\varepsilon\\neq0.\n",
    "$$\n",
    "\n",
    "Think of $a$ as a value and $b$ as \"the derivative riding along\". The rule\n",
    "$\\varepsilon^2=0$ is the *exact* algebraic encoding of \"discard second-order terms\",\n",
    "the very move we used to derive every differentiation rule in\n",
    "that section. Addition is componentwise, and\n",
    "multiplication, dropping the $\\varepsilon^2$ term, is\n",
    "\n",
    "$$\n",
    "(a+b\\varepsilon)(c+d\\varepsilon) = ac + (ad+bc)\\,\\varepsilon.\n",
    "$$\n",
    "\n",
    "Stare at the $\\varepsilon$ coefficient: $ad+bc$ is the *product rule*. This is no\n",
    "coincidence.\n",
    "\n",
    "**Proposition (dual numbers compute derivatives).** *For $f$ expressed as a\n",
    "composition of the primitives below (constants, the identity, sums, products, and\n",
    "elementary functions such as $\\sin$ and $\\exp$), each differentiable at the point\n",
    "where it is evaluated, evaluating that composition on the dual number\n",
    "$x+1\\cdot\\varepsilon$ yields*\n",
    "\n",
    "$$\n",
    "f(x + \\varepsilon) = f(x) + f'(x)\\,\\varepsilon.\n",
    "$$\n",
    "\n",
    "**Proof.** One point of order first: for an elementary primitive such as\n",
    "$\\sin$ or $\\exp$, differentiable at the point where it is evaluated, we *define*\n",
    "its action on a dual argument by\n",
    "\n",
    "$$\n",
    "f(a + b\\varepsilon) := f(a) + f'(a)\\,b\\,\\varepsilon,\n",
    "$$\n",
    "\n",
    "which is exactly what the class methods below implement, so\n",
    "the equation for a single primitive holds *by definition*, not by a\n",
    "Taylor expansion evaluated on a nilpotent argument. (The definition is not\n",
    "arbitrary: for polynomials the algebra alone forces it, since $\\varepsilon^2=0$\n",
    "kills every higher-order term, and the elementary functions adopt the same rule.) The\n",
    "substance of the proposition is that the property survives *assembly*. It holds\n",
    "for constants and the identity ($x+\\varepsilon$ itself), and it is closed under\n",
    "the operations a program is built from. For a sum, $\\varepsilon$-coefficients\n",
    "add, matching $(f+g)'=f'+g'$. For a product, the equation gives\n",
    "coefficient $f'(x)g(x)+f(x)g'(x)$, the product rule. For a composition,\n",
    "substitute $g(x)+g'(x)\\varepsilon$ into a primitive $f$: the definition above,\n",
    "read with $a = g(x)$ and $b = g'(x)$, returns\n",
    "$f(g(x))+f'(g(x))g'(x)\\,\\varepsilon$: the chain rule, and consistent with the\n",
    "sum and product cases. Since $f$ is assembled from these primitives,\n",
    "induction on its expression gives the equation. $\\blacksquare$\n",
    "\n",
    "So differentiation is *free*: run the ordinary computation in the dual-number\n",
    "algebra, set the input's $\\varepsilon$-part to $1$ (the \"seed\"), and the output's\n",
    "$\\varepsilon$-part is the derivative. No formula for $f'$ is ever written down. The\n",
    "implementation is a few lines of operator overloading."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-dual-numbers"
   },
   "source": [
    "class Dual:\n",
    "    \"\"\"A dual number a + b*eps with eps^2 = 0; b carries the derivative.\"\"\"\n",
    "    def __init__(self, a, b=0.0):\n",
    "        self.a, self.b = a, b                       # value, derivative\n",
    "    def __add__(self, o):\n",
    "        o = o if isinstance(o, Dual) else Dual(o)\n",
    "        return Dual(self.a + o.a, self.b + o.b)\n",
    "    __radd__ = __add__\n",
    "    def __mul__(self, o):\n",
    "        o = o if isinstance(o, Dual) else Dual(o)\n",
    "        return Dual(self.a * o.a, self.a * o.b + self.b * o.a)   # product rule\n",
    "    __rmul__ = __mul__\n",
    "    def sin(self):\n",
    "        return Dual(np.sin(self.a), np.cos(self.a) * self.b)     # chain rule\n",
    "    def exp(self):\n",
    "        e = np.exp(self.a)\n",
    "        return Dual(e, e * self.b)\n",
    "    def __repr__(self):\n",
    "        return f'{self.a:.6f} + {self.b:.6f} eps'"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-6 fw=all -->\n",
    "\n",
    "To differentiate $f(x)=\\sin(x^2)+e^x$ at a point, we evaluate it on\n",
    "$x+1\\cdot\\varepsilon$ and read the two parts of the result. The value should match\n",
    "$f$ and the $\\varepsilon$-part should match $f'(x)=2x\\cos(x^2)+e^x$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-dual-eval"
   },
   "source": [
    "def f(x):\n",
    "    return (x * x).sin() + x.exp()                  # sin(x^2) + e^x, in dual algebra\n",
    "\n",
    "x0 = 1.3\n",
    "out = f(Dual(x0, 1.0))                              # seed the derivative with 1\n",
    "exact = 2 * x0 * np.cos(x0**2) + np.exp(x0)         # f'(x) by hand\n",
    "print('dual result :', out)\n",
    "print('value f(x)  :', np.sin(x0**2) + np.exp(x0))\n",
    "print(\"derivative  : dual %.6f  vs  exact %.6f\" % (out.b, exact))"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "dual result : 4.662200 + 3.360101 eps\nvalue f(x)  : 4.662200318713363\nderivative  : dual 3.360101  vs  exact 3.360101\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-7 fw=all -->\n",
    "\n",
    "Value and derivative emerge from a single evaluation. The same machinery handles a\n",
    "vector-valued $\\mathbf f:\\mathbb R^n\\to\\mathbb R^m$ once we let the dual part be a\n",
    "*vector* tangent: seeding the inputs with a direction $\\mathbf v$ and overloading the\n",
    "arithmetic exactly as above propagates the **Jacobian--vector product**\n",
    "$\\mathbf J\\mathbf v$ (a JVP) in the $\\varepsilon$-slot, in a single forward pass and\n",
    "*without ever forming $\\mathbf J$*. Choosing the seed $\\mathbf v=\\mathbf e_j$ then\n",
    "reads off the $j$-th *column* of the Jacobian, so $n$ such passes assemble the whole\n",
    "matrix. Let us close the loop on the earlier $\\mathbb R^2\\to\\mathbb R^2$ map\n",
    "$\\mathbf f(x,y)=(x^2y,\\ \\sin(x+y))$: we run it through the dual algebra with a\n",
    "vector tangent, seed $\\mathbf e_1$ then $\\mathbf e_2$ to recover the two Jacobian\n",
    "columns, and check against the analytic $\\mathbf J(\\mathbf x_0)$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-forward-jvp"
   },
   "source": [
    "class VDual:\n",
    "    \"\"\"Dual number a + b*eps with a *vector* tangent b, so one pass gives J v.\"\"\"\n",
    "    def __init__(self, a, b=0.0):\n",
    "        self.a, self.b = a, np.asarray(b, dtype=float)           # value, tangent\n",
    "    def __add__(self, o):\n",
    "        o = o if isinstance(o, VDual) else VDual(o)\n",
    "        return VDual(self.a + o.a, self.b + o.b)\n",
    "    __radd__ = __add__\n",
    "    def __mul__(self, o):\n",
    "        o = o if isinstance(o, VDual) else VDual(o)\n",
    "        return VDual(self.a * o.a, self.a * o.b + self.b * o.a)   # product rule\n",
    "    __rmul__ = __mul__\n",
    "    def sin(self):\n",
    "        return VDual(np.sin(self.a), np.cos(self.a) * self.b)     # chain rule\n",
    "\n",
    "def f_dual(x, y):\n",
    "    return [(x * x) * y, (x + y).sin()]            # [x^2 y, sin(x+y)] in dual algebra\n",
    "\n",
    "x0 = np.array([1.0, 0.5])\n",
    "cols = []\n",
    "for seed in (np.array([1., 0.]), np.array([0., 1.])):            # e1, then e2\n",
    "    x = VDual(x0[0], seed[0]); y = VDual(x0[1], seed[1])\n",
    "    out = f_dual(x, y)                             # one forward pass per seed\n",
    "    cols.append(np.array([out[0].b, out[1].b]))    # tangent = J @ seed = a column\n",
    "J_fwd = np.column_stack(cols)                      # columns J e1, J e2 assemble J\n",
    "\n",
    "x, y = x0                                          # analytic Jacobian, for the check\n",
    "J_exact = np.array([[2 * x * y, x**2],\n",
    "                    [np.cos(x + y), np.cos(x + y)]])\n",
    "print('forward-mode J:\\n', J_fwd.round(6))\n",
    "print('JVP matches analytic Jacobian:', np.allclose(J_fwd, J_exact))"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward-mode J:\n [[1.       1.      ]\n [0.070737 0.070737]]\nJVP matches analytic Jacobian: True\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-8 fw=all -->\n",
    "\n",
    "Each forward pass produced one column of $\\mathbf J$ in its $\\varepsilon$-slot, with\n",
    "no Jacobian ever materialized; the reverse-mode tape below is verified the same way,\n",
    "against a framework's autograd. So forward mode is cheap for *tall* Jacobians (many\n",
    "outputs, few inputs, $m\\gg n$) and expensive when $n$ is large, exactly the wrong\n",
    "regime for a deep network, whose loss has a single scalar output and millions of\n",
    "inputs. For that we need the other parenthesization.\n",
    "\n",
    "## Reverse-Mode AD, the Tape, and Backprop\n",
    "\n",
    "Recall the associativity observation from the equation: the\n",
    "chain of Jacobians may be multiplied in any order. Forward mode multiplies\n",
    "*right-to-left*, building up $\\mathbf J\\mathbf v$ one input direction at a time.\n",
    "*Reverse mode* multiplies *left-to-right*: it computes the **vector--Jacobian\n",
    "product** $\\mathbf u^\\top\\mathbf J$ (a VJP), one *output* direction at a time.\n",
    "\n",
    "### Why Reverse Mode Is the Right Cost Model\n",
    "\n",
    "The asymmetry is decisive for deep learning. Consider a scalar loss $L:\\mathbb\n",
    "R^n\\to\\mathbb R$ with $n$ parameters. Its Jacobian is a single row ($m=1$), so\n",
    "seeding the reverse pass with $\\mathbf u=1$ at the output computes the *entire*\n",
    "gradient $\\nabla L$ in **one** backward sweep. Forward mode, by contrast, would need\n",
    "$n$ passes, one per parameter direction, to assemble the same row. The cost rule:\n",
    "\n",
    "* **Reverse mode** costs one pass per *output*: cheap for *wide* Jacobians\n",
    "  ($m\\ll n$). A scalar loss is the extreme wide case, so its full gradient costs a\n",
    "  small constant multiple of one forward evaluation (typically $2$--$4\\times$ in\n",
    "  practice; for arithmetic circuits, programs built from $+,-,\\times,\\div$, the\n",
    "  Baur--Strassen theorem pins the multiple at no more than $5$\n",
    "  [@Baur.Strassen.1983]). This is one sweep regardless of how many input\n",
    "  partial derivatives are requested; its absolute work still scales with the\n",
    "  cost of the primal program, which usually grows with $n$.\n",
    "* **Forward mode** costs one pass per *input*: cheap for *tall* Jacobians\n",
    "  ($m\\gg n$).\n",
    "\n",
    "the figure turns the rule into shapes: what a pass\n",
    "buys is one linear combination of rows (reverse) or columns (forward). A\n",
    "coordinate seed selects one literal row or column, so the cheap mode is the one\n",
    "whose number of coordinate seeds matches the Jacobian's short side.\n",
    "\n",
    "![The shape of the Jacobian dictates the cheap mode. A scalar loss has a $1\\times n$ Jacobian: one row, hence one VJP, one backward pass. A one-input map has an $m\\times 1$ Jacobian: one column, one JVP, one forward pass. The full $m\\times n$ matrix costs $\\min(m,n)$ passes either way.](../img/mdl-cal-jacobian-shapes.svg)\n",
    "\n",
    "A network has millions of parameters and one loss, so reverse mode wins by a factor\n",
    "of the parameter count: millions. We will count the passes with our own two\n",
    "engines once both are built. The gradient of a scalar at a small constant multiple\n",
    "of one function evaluation, regardless of parameter count, is the *cheap-gradient\n",
    "principle* [@Baur.Strassen.1983; @Griewank.Walther.2008]; it is what makes\n",
    "training deep networks computationally feasible. **This is backpropagation**\n",
    "[@Linnainmaa.1970; @Speelpenning.1980; @Rumelhart.Hinton.Williams.ea.1988]: the\n",
    "algorithm of that section, where we insisted on\n",
    "\"keeping $\\partial f$ in the numerator\", is reverse-mode AD, and the rule is now\n",
    "rigorous.\n",
    "\n",
    "### The Tape: Record Forward, Replay Backward\n",
    "\n",
    "Reverse mode cannot be done by a value-carrying algebra like dual numbers, because\n",
    "the multiplications run *opposite* to the program's execution: we need each\n",
    "operation's local Jacobian *after* the forward pass has finished. So we *record*\n",
    "the computation as it runs, and then walk the record backward. The record is\n",
    "called the *Wengert list* or *tape*, after the forward-mode paper\n",
    "[@Wengert.1964] that introduced the list of elementary operations; running\n",
    "it *backward* is the later, reverse-mode idea [@Linnainmaa.1970]. The forward pass builds a directed acyclic graph of\n",
    "elementary operations and stashes the values needed for each local derivative; the\n",
    "backward pass seeds the output *adjoint* (the running $\\partial L/\\partial\\cdot$) and\n",
    "pushes it through each recorded operation's VJP, accumulating contributions where a\n",
    "value feeds several consumers. the figure contrasts the two\n",
    "sweeps.\n",
    "\n",
    "![Forward-mode versus reverse-mode automatic differentiation on a computation graph. Forward mode (top) propagates a Jacobian--vector product left to right, in lock-step with evaluation; a coordinate seed selects one Jacobian column. Reverse mode (bottom) records a tape on a forward pass, then propagates a vector--Jacobian product right to left; a coordinate output seed selects one row, and a scalar output yields its full gradient.](../img/mdl-cal-fwd-vs-rev.svg)\n",
    "\n",
    "A tiny tape makes this concrete. We record each operation as a node holding its\n",
    "value and a *backward* closure that, given the adjoint of the node's output,\n",
    "increments the adjoints of its inputs by the appropriate VJP. Replaying the tape in\n",
    "reverse and seeding the final adjoint with $1$ leaves $\\partial L/\\partial\\cdot$ in\n",
    "every input's adjoint slot."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-tape"
   },
   "source": [
    "class Var:\n",
    "    \"\"\"A scalar node on the autodiff tape; .grad accumulates the adjoint.\"\"\"\n",
    "    def __init__(self, value, parents=(), backward=lambda g: None):\n",
    "        self.value = value\n",
    "        self.grad = 0.0\n",
    "        self._parents = parents                     # nodes this one depends on\n",
    "        self._backward = backward                   # pushes adjoint to parents\n",
    "    def __add__(self, o):\n",
    "        o = o if isinstance(o, Var) else Var(o)\n",
    "        out = Var(self.value + o.value, (self, o))\n",
    "        def back(g):                                # d(a+b) = 1*g to each parent\n",
    "            self.grad += g; o.grad += g\n",
    "        out._backward = back\n",
    "        return out\n",
    "    __radd__ = __add__\n",
    "    def __mul__(self, o):\n",
    "        o = o if isinstance(o, Var) else Var(o)\n",
    "        out = Var(self.value * o.value, (self, o))\n",
    "        def back(g):                                # product rule, reversed\n",
    "            self.grad += o.value * g; o.grad += self.value * g\n",
    "        out._backward = back\n",
    "        return out\n",
    "    __rmul__ = __mul__\n",
    "\n",
    "def backprop(node):\n",
    "    \"\"\"Topologically sort the tape and replay it backward from `node`.\"\"\"\n",
    "    order, seen = [], set()\n",
    "    def visit(v):\n",
    "        if v not in seen:\n",
    "            seen.add(v)\n",
    "            for p in v._parents:\n",
    "                visit(p)\n",
    "            order.append(v)\n",
    "    visit(node)\n",
    "    node.grad = 1.0                                 # seed the output adjoint\n",
    "    for v in reversed(order):\n",
    "        v._backward(v.grad)"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-9 fw=all -->\n",
    "\n",
    "We differentiate an expression with the same diamond-reuse structure that made\n",
    "the backward sweep of that section efficient, and\n",
    "check the tape's gradients against the framework's autograd. Take $g(u,v)=r^2$ with the shared\n",
    "intermediate $r=uv+u$ (in code, `r = u*v + u` followed by `y = r*r`). One forward\n",
    "pass records the tape, one backward pass yields both partials. Its graph,\n",
    "the figure, is not a chain but a *diamond*: $u$ feeds both the\n",
    "product $uv$ and the sum $r$, and $r$ feeds *both* arguments of the final\n",
    "multiplication. This reuse is exactly why the backward pass must visit nodes in\n",
    "reverse topological order (each node before the nodes it depends on, so that its\n",
    "adjoint is complete before it is pushed onward) and *accumulate* adjoints with\n",
    "`+=` rather than\n",
    "overwrite them: when the node $y=r\\cdot r$ runs its backward step, it pushes the\n",
    "contribution $\\bar y\\,r$ into `r.grad` *twice*, once per argument, so the adjoint\n",
    "$\\bar r = 2r\\bar y$ arrives as a sum: a node's adjoint is the *sum* of the\n",
    "contributions along every outgoing edge, the chain rule's \"sum over paths\" made\n",
    "operational.\n",
    "\n",
    "![The computation graph for the tape example $g(u,v)=r^2$ with $r=uv+u$. It is a diamond, not a chain: $u$ fans out to the product $t=uv$ and the sum $r=t+u$, and $r$ feeds both arguments of $y=r\\cdot r$ (the doubled edge). Because a value can feed several consumers, the backward pass walks the tape in reverse topological order and accumulates each adjoint over outgoing edges with `+=`; here `r.grad` receives the contribution $\\bar y\\,r$ twice, summing to $\\bar r=2r\\bar y$.](../img/mdl-cal-tape-dag.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "matrix-calculus-autodiff-tape-check"
   },
   "source": [
    "u, v = Var(2.0), Var(-3.0)\n",
    "r = u * v + u                                       # shared intermediate\n",
    "y = r * r                                           # r feeds both arguments\n",
    "backprop(y)\n",
    "def g(u, v):\n",
    "    r = u * v + u\n",
    "    return r * r\n",
    "du, dv = jax.grad(g, argnums=(0, 1))(2.0, -3.0)     # framework autograd\n",
    "print('our tape   : dy/du = %.4f  dy/dv = %.4f' % (u.grad, v.grad))\n",
    "print('jax        : dy/du = %.4f  dy/dv = %.4f' % (float(du), float(dv)))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "our tape   : dy/du = 16.0000  dy/dv = -16.0000\njax        : dy/du = 16.0000  dy/dv = -16.0000\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-10 fw=all -->\n",
    "\n",
    "Our thirty-line tape reproduces the framework's gradients exactly. The real systems\n",
    "generalize it in three ways that do not change the principle: the nodes hold\n",
    "*tensors* rather than scalars, so each `_backward` is a VJP (a matrix operation, e.g.\n",
    "the $\\mathbf x^\\top$ in the equation) rather than a scalar multiply;\n",
    "the primitive set is large (every differentiable op ships a VJP rule); and the graph\n",
    "is built either eagerly (PyTorch's dynamic tape) or by tracing (JAX, TensorFlow).\n",
    "But the skeleton (record forward, seed the output adjoint, replay backward\n",
    "accumulating VJPs) is exactly what you wrote above and exactly what\n",
    "the backward pass does.\n",
    "\n",
    "### The Cost Asymmetry, Counted\n",
    "\n",
    "With both engines on the table, the cost claim that opened this half of the\n",
    "section stops being an assertion. Take a scalar function of $n=200$ inputs,\n",
    "built from nothing but the additions and multiplications both engines support.\n",
    "Forward mode must run one pass per input direction to assemble the gradient (a\n",
    "$1\\times n$ Jacobian, column by column); the tape runs *one* backward sweep."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-matrix-calculus-autodiff-the-cost-asymmetry-measured"
   },
   "source": [
    "n = 200\n",
    "def f(xs):                       # f: R^n -> R, from + and * alone\n",
    "    total = 0.0\n",
    "    for x in xs:\n",
    "        total = total + x * x\n",
    "    return total * total         # (sum of squares)^2\n",
    "\n",
    "xs0 = [0.01 * (j + 1) for j in range(n)]\n",
    "passes, grad_fwd = 0, []\n",
    "for j in range(n):               # forward mode: one pass per input direction\n",
    "    duals = [Dual(x, 1.0 if k == j else 0.0) for k, x in enumerate(xs0)]\n",
    "    grad_fwd.append(f(duals).b)\n",
    "    passes += 1\n",
    "\n",
    "xs = [Var(x) for x in xs0]       # reverse mode: one recorded forward pass...\n",
    "backprop(f(xs))                  # ...and one backward sweep for all n partials\n",
    "grad_rev = [x.grad for x in xs]\n",
    "\n",
    "gap = max(abs(a - b) for a, b in zip(grad_fwd, grad_rev))\n",
    "print(f'forward mode : {passes} passes (one per input)')\n",
    "print(f'reverse mode : 1 backward sweep for all {n} partials')\n",
    "print(f'max |forward - reverse| over the gradient: {gap:.1e}')"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward mode : 200 passes (one per input)\nreverse mode : 1 backward sweep for all 200 partials\nmax |forward - reverse| over the gradient: 0.0e+00\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-11 fw=all -->\n",
    "\n",
    "The two gradients agree to machine precision, but the pass counts differ by a\n",
    "factor of $n$: $200$ forward passes against one backward sweep, which is the\n",
    "loop structure itself. Each forward-mode pass costs about one evaluation of $f$,\n",
    "and the tape's recorded forward pass plus backward sweep costs a small constant\n",
    "multiple of one, so the *work* ratio is $\\Theta(n)$, i.e. proportional to $n$ up\n",
    "to a constant factor. Nothing about the ratio is special to $200$: replace it by\n",
    "the $10^8$ parameters of a modern network and this factor is the \"factor of\n",
    "millions\" claimed above, which is why no framework computes training gradients\n",
    "in forward mode.\n",
    "\n",
    "### Never Form the Jacobian\n",
    "\n",
    "The two engines above expose exactly two products, the forward-mode JVP\n",
    "$\\mathbf J\\mathbf v$ and the reverse-mode VJP $\\mathbf u^\\top\\mathbf J$, and every\n",
    "quantity a training loop needs is built by *composing* those, never by materializing\n",
    "$\\mathbf J$. This is deliberate. Writing down a dense $m\\times n$ Jacobian costs\n",
    "$\\min(m,n)$ passes (one per column via JVP, or one per row via VJP) plus $\\Theta(mn)$\n",
    "storage, while the thing you actually want (a gradient $\\mathbf u^\\top\\mathbf J$, a\n",
    "directional derivative $\\mathbf J\\mathbf v$, a sequence of these through a deep\n",
    "composition) costs a single pass and never holds more than a vector. Autograd\n",
    "systems therefore expose the JVP and the VJP as their primitives and treat the\n",
    "full Jacobian as a rare, opt-in convenience. Even that convenience is built from\n",
    "the same parts: a Jacobian, when one is truly needed, is assembled by batching\n",
    "VJPs over the $m$ output directions (rows) or JVPs over the $n$ input directions\n",
    "(columns), and a Hessian by the forward-over-reverse composition of the next\n",
    "subsection, when $n$ is small enough to afford one. If you find\n",
    "yourself assembling a Jacobian, you have almost certainly written down a matrix you\n",
    "could have multiplied through.\n",
    "\n",
    "### Hessian-Vector Products: One Order Up\n",
    "\n",
    "The same \"don't materialize the matrix\" discipline pays off a second time, one\n",
    "derivative higher. Second-order optimizers and curvature diagnostics never need\n",
    "the full Hessian $\\mathbf H=\\nabla^2 L$, an $n\\times n$ object, hopeless to even\n",
    "store for $n$ in the millions. Newton's method, the *conjugate-gradient* (CG)\n",
    "method (an iterative linear-system solver that touches its matrix only through\n",
    "matrix-vector products), and *Gauss--Newton* (which approximates the Hessian by\n",
    "a product of first-derivative Jacobians) [@Nocedal.Wright.2006], see\n",
    "that section, all need only\n",
    "the **Hessian--vector product** $\\mathbf H\\mathbf v$: the action of curvature in a\n",
    "single direction, e.g. one matrix-vector multiply inside a CG iteration. Pearlmutter's\n",
    "trick [@Pearlmutter.1994] obtains it *without forming $\\mathbf H$* by composing\n",
    "the two modes. Since $\\mathbf H = \\nabla^2 L$ is the Jacobian of the gradient map\n",
    "$\\nabla L$, the Hessian--vector product is a *directional derivative of the gradient*,\n",
    "\n",
    "$$\n",
    "\\mathbf H\\mathbf v\n",
    "   = \\mathbf J_{\\nabla L}\\,\\mathbf v\n",
    "   = \\left.\\frac{d}{ds}\\right|_{s=0}\\nabla L(\\mathbf x + s\\mathbf v),\n",
    "$$\n",
    "\n",
    "so we push the tangent $\\mathbf v$ through the gradient computation in *forward over\n",
    "reverse*: a forward-mode JVP applied to the reverse-mode gradient. The cost is a small\n",
    "constant multiple of one gradient evaluation: curvature in one chosen direction\n",
    "without $n$ separate gradient evaluations [@Baydin.Pearlmutter.Radul.ea.2018].\n",
    "Its absolute work still scales with the computation of $L$ and therefore usually\n",
    "with the problem size.\n",
    "In terms of our toy engines, layering a `Dual` seed over the reverse-mode `Var`\n",
    "tape is exactly this forward-over-reverse construction. It is\n",
    "what makes Newton and conjugate-gradient methods\n",
    "and the curvature analysis of that section tractable at scale.\n",
    "\n",
    "The recipe is short enough to show whole. For the quadratic\n",
    "$L(\\mathbf x)=\\tfrac12\\mathbf x^\\top\\mathbf A\\mathbf x$ with symmetric\n",
    "$\\mathbf A$, the gradient is $\\mathbf A\\mathbf x$ and the Hessian is $\\mathbf A$\n",
    "itself, so the correct answer is known in advance:\n",
    "$\\mathbf H\\mathbf v = \\mathbf A\\mathbf v$. No $2\\times2$ matrix of second\n",
    "derivatives is ever assembled below; only the gradient map appears, differentiated\n",
    "once in the direction $\\mathbf v$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-matrix-calculus-autodiff-hessian-vector-products-one-order-up"
   },
   "source": [
    "# Hv for L(x) = x^T A x / 2, where H = A exactly\n",
    "A = jnp.array([[3.0, 1.0], [1.0, 2.0]])\n",
    "L = lambda x: 0.5 * x @ A @ x\n",
    "x, v = jnp.array([1.0, -2.0]), jnp.array([1.0, 0.5])\n",
    "_, Hv = jax.jvp(jax.grad(L), (x,), (v,))      # forward over reverse\n",
    "print('forward-over-reverse Hv:', Hv)\n",
    "print('A v                    :', A @ v)"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward-over-reverse Hv: [3.5 2. ]\nA v                    : [3.5 2. ]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-12 fw=all -->\n",
    "\n",
    "The printed product equals $\\mathbf A\\mathbf v$ to the digit. Exercise 7 works\n",
    "through the analytic claim and contrasts forward-over-reverse with the\n",
    "*double-backward* (reverse-over-reverse) alternative.\n",
    "\n",
    "### Differentiating Through Equations\n",
    "\n",
    "Everything so far differentiated an *explicit* program: a chain of primitives we\n",
    "can record on a tape. But some of the most useful outputs in machine learning\n",
    "are defined *implicitly*, as the solution of an equation: the fixed point a\n",
    "solver converges to, the equilibrium of a physical model, the argmin of an inner\n",
    "optimization. Suppose $\\mathbf x^\\star(\\boldsymbol\\theta)$ is defined by the\n",
    "condition\n",
    "\n",
    "$$\n",
    "\\mathbf g\\bigl(\\mathbf x^\\star(\\boldsymbol\\theta),\\,\\boldsymbol\\theta\\bigr) = \\mathbf 0,\n",
    "$$\n",
    "\n",
    "and we want $\\partial\\mathbf x^\\star/\\partial\\boldsymbol\\theta$, say to train\n",
    "$\\boldsymbol\\theta$ by gradient descent through the solver. Taping the solver's\n",
    "dozens of iterations works but is wasteful and fragile. The *implicit function\n",
    "theorem* [@Rudin.1976], which we use without proof, says we never have to:\n",
    "if $\\mathbf g$ is continuously differentiable\n",
    "and the Jacobian $\\partial\\mathbf g/\\partial\\mathbf x$ is invertible at the\n",
    "solution, then $\\mathbf x^\\star(\\boldsymbol\\theta)$ is locally a differentiable\n",
    "function of $\\boldsymbol\\theta$, and differentiating the identity\n",
    "$\\mathbf g(\\mathbf x^\\star(\\boldsymbol\\theta),\\boldsymbol\\theta)\\equiv\\mathbf 0$\n",
    "by the chain rule the equation gives\n",
    "$\\frac{\\partial\\mathbf g}{\\partial\\mathbf x}\\frac{\\partial\\mathbf x^\\star}{\\partial\\boldsymbol\\theta}\n",
    "+\\frac{\\partial\\mathbf g}{\\partial\\boldsymbol\\theta}=\\mathbf 0$, hence\n",
    "\n",
    "$$\n",
    "\\frac{\\partial\\mathbf x^\\star}{\\partial\\boldsymbol\\theta}\n",
    "   = -\\Bigl(\\frac{\\partial\\mathbf g}{\\partial\\mathbf x}\\Bigr)^{\\!-1}\n",
    "     \\frac{\\partial\\mathbf g}{\\partial\\boldsymbol\\theta}.\n",
    "$$\n",
    "\n",
    "Both Jacobians on the right are evaluated *at the solution only*: the gradient\n",
    "of an implicitly defined output costs one linear solve, no matter how many\n",
    "iterations the solver ran. The \"never form the Jacobian\" discipline applies\n",
    "verbatim: a VJP $\\mathbf u^\\top\\partial\\mathbf x^\\star/\\partial\\boldsymbol\\theta$\n",
    "needs only one linear solve against\n",
    "$(\\partial\\mathbf g/\\partial\\mathbf x)^\\top$, which is exactly the *adjoint\n",
    "method* that that section uses to backpropagate through ODE\n",
    "solvers. Deep equilibrium models [@Bai.Kolter.Koltun.2019] make this the\n",
    "whole network: a layer's output is the fixed point\n",
    "$\\mathbf x^\\star=\\mathbf f(\\mathbf x^\\star,\\boldsymbol\\theta)$, trained through\n",
    "the equation alone, and bilevel problems such as hyperparameter\n",
    "optimization differentiate through an inner argmin the same way.\n",
    "The next cell verifies the equation on a $2\\times2$ nonlinear\n",
    "system, against a finite difference of the solver's output."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-matrix-calculus-autodiff-differentiating-through-equations"
   },
   "source": [
    "# Implicit gradient dx*/dtheta = -(dg/dx)^{-1} dg/dtheta on a 2x2 system,\n",
    "# checked against a finite difference of the solver output\n",
    "def g(x, th):                        # g(x, theta) = 0 defines x*(theta)\n",
    "    return np.array([x[0]**2 + th * x[1] - 2, x[0] * x[1] + th - 1])\n",
    "\n",
    "def dg_dx(x, th):                    # Jacobian in x, evaluated at (x, th)\n",
    "    return np.array([[2 * x[0], th], [x[1], x[0]]])\n",
    "\n",
    "def solve(th):                       # Newton iteration, from a fixed start\n",
    "    x = np.array([1.0, 1.0])\n",
    "    for _ in range(20):\n",
    "        x = x - np.linalg.solve(dg_dx(x, th), g(x, th))\n",
    "    return x\n",
    "\n",
    "th, h = 0.3, 1e-6\n",
    "xs = solve(th)\n",
    "dg_dth = np.array([xs[1], 1.0])      # dg/dtheta at the solution\n",
    "implicit = -np.linalg.solve(dg_dx(xs, th), dg_dth)\n",
    "fd = (solve(th + h) - solve(th - h)) / (2 * h)\n",
    "print('implicit-function gradient:', implicit.round(8))\n",
    "print('finite-difference check   :', fd.round(8))"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "implicit-function gradient: [-0.11311451 -0.69322129]\nfinite-difference check   : [-0.11311451 -0.69322129]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-13 fw=all -->\n",
    "\n",
    "The two vectors agree to the precision the finite difference allows. Note what\n",
    "the implicit route never touched: the twenty Newton iterations inside `solve`.\n",
    "Only the residual's two Jacobians at the converged point enter\n",
    "the equation: the equation, not the algorithm that solved it,\n",
    "carries the derivative.\n",
    "\n",
    "### The Memory Trade-off and Checkpointing\n",
    "\n",
    "One cost is now visible. Reverse mode must *store the forward intermediates* (the\n",
    "`value`s on the tape) until the backward pass consumes them, so its memory grows\n",
    "*linearly* with the length $L$ of the computation, whereas forward mode keeps only\n",
    "the current value and tangent. For a deep network this $O(L)$ activation memory is\n",
    "often the binding constraint on batch size and depth. *Gradient checkpointing*\n",
    "[@Chen.Xu.Zhang.ea.2016] buys it back: store only $O(\\sqrt L)$ of the\n",
    "intermediates and *recompute* the rest on the fly during the backward pass, turning\n",
    "$O(L)$ memory into $O(\\sqrt L)$ at the cost of roughly one extra forward\n",
    "pass; the figure draws the schedule. The\n",
    "trade-off can be pushed much further: Griewank's *treeverse* algorithm\n",
    "[@Griewank.1992] checkpoints recursively to reach $O(\\log L)$ memory, at the\n",
    "price of $O(L\\log L)$ recomputation, provably optimal among checkpointing\n",
    "schedules with the given snapshot budget. Forward and reverse modes,\n",
    "mixed modes, and checkpointing are treated in full in Griewank and Walther's\n",
    "monograph *Evaluating Derivatives* [@Griewank.Walther.2008] and in the survey\n",
    "[@Baydin.Pearlmutter.Radul.ea.2018]; the same reverse-mode tape returns as the\n",
    "*adjoint method* for differentiating through ODE solvers in\n",
    "that section.\n",
    "\n",
    "![Gradient checkpointing as a timeline. Plain reverse mode (top) keeps all $L$ forward activations alive until the backward sweep consumes them. Checkpointing (bottom) stores only every $\\sqrt{L}$-th activation; when the backward sweep reaches a segment, that segment is recomputed forward from its checkpoint and then swept: $O(\\sqrt{L})$ memory for roughly one extra forward pass.](../img/mdl-cal-checkpointing.svg)\n",
    "\n",
    "## Summary\n",
    "\n",
    "* The **Jacobian** $\\mathbf J_{\\mathbf f}\\in\\mathbb R^{m\\times n}$ is the best local\n",
    "  linear approximation of a vector-valued $\\mathbf f$:\n",
    "  $\\mathbf f(\\mathbf x+\\boldsymbol\\delta)\\approx\\mathbf f(\\mathbf x)+\\mathbf J\\boldsymbol\\delta$,\n",
    "  with $[\\mathbf J]_{ij}=\\partial f_i/\\partial x_j$. The row gradient ($m=1$) and the\n",
    "  Hessian (Jacobian of $\\nabla f$) are special cases.\n",
    "* The chain rule in matrix form is **Jacobian composition**\n",
    "  $\\mathbf J_{\\mathbf g\\circ\\mathbf f}=\\mathbf J_{\\mathbf g}\\mathbf J_{\\mathbf f}$, so\n",
    "  a depth-$L$ network's derivative is one long matrix product\n",
    "  $\\mathbf J_L\\cdots\\mathbf J_1$. Its *associativity* is the whole story of AD.\n",
    "* **Layout conventions** (numerator vs denominator) differ by a transpose; we use\n",
    "  numerator layout for Jacobians and denominator (column) layout for gradients.\n",
    "* The recurring DL identities all follow from differentiating one component and\n",
    "  reassembling, not from a table: $\\nabla\\,\\mathbf a^\\top\\mathbf x=\\mathbf a$,\n",
    "  $\\nabla\\,\\mathbf x^\\top\\mathbf A\\mathbf x=(\\mathbf A+\\mathbf A^\\top)\\mathbf x$,\n",
    "  the rank-one $\\nabla_{\\mathbf W}\\|\\mathbf W\\mathbf x-\\mathbf y\\|^2=2(\\mathbf W\\mathbf x-\\mathbf y)\\mathbf x^\\top$,\n",
    "  and the softmax--cross-entropy logit gradient $\\mathbf p-\\mathbf y$.\n",
    "* **Forward-mode AD** carries a derivative via *dual numbers* ($\\varepsilon^2=0$);\n",
    "  one pass computes a Jacobian--vector product, a linear combination of columns\n",
    "  (one literal column for a coordinate seed), so it is cheap for tall Jacobians.\n",
    "* **Reverse-mode AD** records a *tape* and replays it backward, computing a\n",
    "  vector--Jacobian product, a linear combination of rows (one literal row for a\n",
    "  coordinate seed). Because a loss is scalar, **backpropagation is reverse-mode\n",
    "  AD**: one backward sweep yields the gradient with respect to every parameter at\n",
    "  a small constant multiple of one forward evaluation, at the price of storing\n",
    "  the forward intermediates. The number of sweeps does not grow with $n$; the\n",
    "  cost of each sweep still follows the cost of the primal computation.\n",
    "* **Never form the Jacobian.** Autograd exposes the JVP and the VJP and you\n",
    "  *compose* them; a dense $m\\times n$ Jacobian costs $\\min(m,n)$ passes and is\n",
    "  almost always avoidable. One order up, the **Hessian--vector product**\n",
    "  $\\mathbf H\\mathbf v$ comes from *forward-over-reverse* (Pearlmutter's trick) at\n",
    "  the cost of one extra gradient: curvature without forming $\\mathbf H$, which is\n",
    "  what makes Newton/CG methods scale. The $O(L)$ activation memory of the tape can be traded down to\n",
    "  $O(\\sqrt L)$ by *gradient checkpointing* at roughly one extra forward pass.\n",
    "* **Implicitly defined outputs** need no tape through the solver: if\n",
    "  $\\mathbf g(\\mathbf x^\\star,\\boldsymbol\\theta)=\\mathbf 0$, the implicit function\n",
    "  theorem gives $\\partial\\mathbf x^\\star/\\partial\\boldsymbol\\theta\n",
    "  =-(\\partial\\mathbf g/\\partial\\mathbf x)^{-1}\\partial\\mathbf g/\\partial\\boldsymbol\\theta$\n",
    "  from the Jacobians at the solution alone: one linear solve per gradient, the\n",
    "  seed of the adjoint method, deep equilibrium models, and bilevel optimization.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Compute the Jacobian of $\\mathbf f(x,y)=(x^2y,\\ \\sin(x+y))$ and verify the linear\n",
    "   approximation $\\mathbf f(\\mathbf x+\\boldsymbol\\delta)\\approx\\mathbf f(\\mathbf x)+\\mathbf J\\boldsymbol\\delta$\n",
    "   at $(1,0)$ against a finite difference. By how much does the error shrink when you\n",
    "   halve $\\boldsymbol\\delta$?\n",
    "2. Show $\\nabla_{\\mathbf x}\\|\\mathbf A\\mathbf x-\\mathbf b\\|_2^2=2\\mathbf A^\\top(\\mathbf A\\mathbf x-\\mathbf b)$\n",
    "   two ways: by an index/Einstein-notation derivation, and by the scalar-collapse\n",
    "   heuristic (guess from $\\frac{d}{dx}(ax-b)^2$, then fix the shapes with transposes).\n",
    "3. For a scalar loss with $n=10^6$ parameters, how many passes does forward mode need\n",
    "   to assemble the full gradient, and how many does reverse mode need? Explain via JVPs\n",
    "   versus VJPs.\n",
    "4. Extend the `Dual` class with a `__pow__` (or a `log`) method and use it to\n",
    "   differentiate $f(x)=\\log(1+x^2)$ at $x=2$; check against the analytic derivative.\n",
    "   Which rule from the proof of the equation does each new method encode?\n",
    "5. Add a new primitive (e.g. $\\exp$ or $\\sin$) to the reverse-mode `Var`/`backprop`\n",
    "   tape, give it the correct `_backward`, and verify a gradient against a framework's\n",
    "   autograd.\n",
    "6. Derive the softmax Jacobian the equation from the quotient\n",
    "   rule, then re-derive the logit gradient $\\mathbf p-\\mathbf y$ and explain why\n",
    "   fusing softmax with cross-entropy is numerically preferable to composing them.\n",
    "7. **Hessian--vector product.** For $L(\\mathbf x)=\\tfrac12\\mathbf x^\\top\\mathbf A\\mathbf x$\n",
    "   with symmetric $\\mathbf A$, show analytically that $\\nabla L=\\mathbf A\\mathbf x$ and\n",
    "   hence $\\mathbf H\\mathbf v=\\mathbf A\\mathbf v$. Then implement the forward-over-reverse\n",
    "   recipe of the equation for a general scalar $L$ using a framework's\n",
    "   autograd (a forward-mode JVP pushed through the reverse-mode gradient: in PyTorch,\n",
    "   `torch.func.jvp(torch.func.grad(L), (x,), (v,))`; in JAX,\n",
    "   `jax.jvp(jax.grad(L), (x,), (v,))`) and confirm it returns $\\mathbf A\\mathbf v$\n",
    "   *without ever building the $n\\times n$ matrix $\\mathbf H$*. PyTorch's classic\n",
    "   *double-backward* alternative (differentiate $\\langle\\nabla L,\\mathbf v\\rangle$\n",
    "   once more, via `torch.autograd.grad(L, x, create_graph=True)`) computes the same\n",
    "   product as *reverse-over-reverse*; it works, but it tapes the backward pass\n",
    "   itself, so it costs somewhat more time and memory than forward-over-reverse.\n",
    "8. **Gradient of the log-absolute-determinant.** For a real invertible\n",
    "   $\\mathbf A$, show that\n",
    "   $\\nabla_{\\mathbf A}\\log|\\det\\mathbf A|=\\mathbf A^{-\\top}$. (If\n",
    "   $\\det\\mathbf A>0$, this is also the gradient of $\\log\\det\\mathbf A$.) *Hint:* perturb one entry\n",
    "   at a time: with $\\mathbf E=\\mathbf e_a\\mathbf e_b^\\top$, expand\n",
    "   $\\det(\\mathbf A+\\varepsilon\\mathbf E)\n",
    "   =\\det(\\mathbf A)\\det(\\mathbf I+\\varepsilon\\mathbf A^{-1}\\mathbf E)\n",
    "   \\approx\\det(\\mathbf A)\\,\\bigl(1+\\varepsilon\\operatorname{tr}(\\mathbf A^{-1}\\mathbf E)\\bigr)$\n",
    "   and read off the $(a,b)$ partial; equivalently, apply Jacobi's formula\n",
    "   $\\tfrac{d}{dt}\\det\\mathbf A(t)=\\det\\mathbf A(t)\\,\n",
    "   \\operatorname{tr}\\bigl(\\mathbf A(t)^{-1}\\tfrac{d\\mathbf A}{dt}\\bigr)$. This\n",
    "   identity is what makes the $\\log|\\det\\mathbf J|$ term of the change-of-variables\n",
    "   formula for densities (that section) trainable by\n",
    "   gradient descent; normalizing flows, generative models built from invertible\n",
    "   maps and trained on exactly this log-likelihood (that section),\n",
    "   depend on it. Verify it numerically by\n",
    "   differentiating the log-absolute-determinant of a random invertible\n",
    "   $3\\times3$ matrix (use a `slogdet`-style routine) with a framework's autograd\n",
    "   and comparing against $\\mathbf A^{-\\top}$.\n",
    "9. **Attention Jacobians.** In self-attention (that section),\n",
    "   softmax is applied *row-wise* to the score matrix\n",
    "   $\\mathbf S=\\mathbf Q\\mathbf K^\\top/\\sqrt d$, giving attention weights\n",
    "   $\\mathbf P$ with rows $\\mathbf p_i=\\operatorname{softmax}(\\mathbf s_i)$.\n",
    "   Show that the Jacobian of the attention weights with respect to the scores is\n",
    "   *block-diagonal*: row $i$ of $\\mathbf P$ depends only on row $i$ of\n",
    "   $\\mathbf S$, and each block is exactly the per-row softmax Jacobian\n",
    "   $\\operatorname{diag}(\\mathbf p_i)-\\mathbf p_i\\mathbf p_i^\\top$ of\n",
    "   the equation. Then push one step further down the chain:\n",
    "   differentiating the scores $\\mathbf s_i=\\mathbf K\\mathbf q_i/\\sqrt d$ with\n",
    "   respect to the query $\\mathbf q_i$ contributes the Jacobian\n",
    "   $\\mathbf K/\\sqrt d$, so\n",
    "   $\\partial\\mathbf p_i/\\partial\\mathbf q_i\n",
    "   =\\bigl(\\operatorname{diag}(\\mathbf p_i)-\\mathbf p_i\\mathbf p_i^\\top\\bigr)\\mathbf K/\\sqrt d$.\n",
    "   Where does the temperature-like $1/\\sqrt d$ end up in the gradient, and what\n",
    "   does it do to gradient magnitudes as $d$ grows? Verify the block-diagonal\n",
    "   structure numerically with a framework's `jacobian` on a $3\\times3$ score\n",
    "   matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-matrix-calculus-autodiff-md-14 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
