{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-1 fw=all -->\n",
    "\n",
    "# Memory and Precision\n",
    "\n",
    "So far this chapter has fought for *time*: fewer memory round trips, fewer\n",
    "kernel launches, more arithmetic per byte. This section fights for *space*.\n",
    "Sooner or later a model stops fitting: the optimizer states, the\n",
    "activations held for the backward pass, and a large enough batch together\n",
    "overflow the card, and training dies with an out-of-memory error long\n",
    "before it is slow. The good news is that memory is the most *legible* of\n",
    "the resources — you can write down exactly where every byte of a training\n",
    "step goes, measure it against that prediction, and then spend three\n",
    "well-understood techniques to buy headroom: lower-precision formats, which\n",
    "also buy speed; activation checkpointing, which trades compute for memory;\n",
    "and gradient accumulation, which trades *time* for the effect of a bigger\n",
    "batch. Together these are the \"it doesn't fit\" rung of the ladder, and\n",
    "they are what stands between a 24 GB card and a model that wants 30.\n",
    "\n",
    "*Prerequisites: the format ladder and tensor-core requirement of*\n",
    "that section*; the* `d2l.Benchmark` *timer of*\n",
    "that section*; the AdamW optimizer states of*\n",
    "that section*. This section finally discharges the \"gradient\n",
    "accumulation, returned to in the performance chapter\" promise of*\n",
    "that section*.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-memory-and-precision"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-2 fw=all -->\n",
    "\n",
    "## The Memory Anatomy of a Training Step\n",
    "\n",
    "A training step holds four kinds of tensor in memory, and knowing their\n",
    "sizes lets you predict the peak before you run anything\n",
    "(the figure). Let the model have $P$ parameters.\n",
    "\n",
    "* **Parameters** — $P$ values. In mixed-precision training these are kept\n",
    "  in fp32 (4 bytes) as the master copy: $4P$ bytes.\n",
    "* **Gradients** — one per parameter, produced during the backward pass:\n",
    "  another $4P$ bytes.\n",
    "* **Optimizer states** — Adam keeps a first and second moment per\n",
    "  parameter (that section), $2 \\times 4P = 8P$ bytes; plain SGD\n",
    "  keeps none.\n",
    "* **Activations** — every intermediate result the forward pass produces\n",
    "  that the backward pass will need. Unlike the first three, this scales\n",
    "  with the *batch* and the *depth*: roughly\n",
    "  $\\textrm{batch} \\times \\textrm{sequence} \\times \\textrm{width} \\times\n",
    "  \\textrm{depth}$, and it is usually the term that blows up.\n",
    "\n",
    "![Where a training step's memory goes, over the step timeline. Parameters\n",
    "and optimizer states are constant; activations grow through the forward\n",
    "pass and are released through the backward pass; gradients fill in as\n",
    "they are computed. Peak memory lands early in the backward pass, when\n",
    "activations are still full.](../img/mdl-perf-memory-anatomy.svg)\n",
    "\n",
    "Add the constant terms: training a model with Adam costs about\n",
    "$4P + 4P + 8P = 16P$ bytes *before a single activation* — sixteen bytes\n",
    "per parameter, the number to carry in your head. (Mixed precision, below,\n",
    "nudges this to roughly 18–20 by keeping some tensors in two formats; it\n",
    "still saves memory overall, because activations, the dominant term,\n",
    "halve.) The shape of the figure is the other half of the\n",
    "lesson: activations *accumulate* through the forward pass and are *freed*\n",
    "as the backward pass consumes them, so peak memory occurs early in\n",
    "backward — which is exactly the moment activation checkpointing will\n",
    "target.\n",
    "\n",
    "## Measuring Memory\n",
    "\n",
    "The frameworks let you check the anatomy against reality, and — as with\n",
    "timing — the two do it in characteristically different ways. PyTorch keeps\n",
    "running *counters* you query at runtime; JAX lets the compiler *plan*\n",
    "memory ahead of time and reports the plan. Start with PyTorch's counters,\n",
    "the numbers the anatomy table predicts:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-3 fw=all -->\n",
    "\n",
    "The measured peak overshoots the $16P$ prediction by roughly 30%, and the\n",
    "gap has an anatomy of its own: this stack's activations (about 13 MB),\n",
    "autograd bookkeeping and cuBLAS workspaces, temporaries that the fused\n",
    "optimizer update allocates (about the size of one more gradient set), and\n",
    "the caching allocator's block rounding. The second printed number is\n",
    "larger again: `max_memory_allocated` counts bytes inside live tensors,\n",
    "while `max_memory_reserved` is the caching allocator's high-water mark —\n",
    "the distinction that section dissects. So read $16P$ as this\n",
    "*configuration's* floor — fp32 parameters under Adam — rather than as the\n",
    "whole anatomy; and whether activations dominate the total, as they do in\n",
    "the transformer regime, is workload-dependent: batch, sequence length,\n",
    "and depth decide.\n",
    "\n",
    "For a finer picture, PyTorch can record every allocation and free as a\n",
    "*snapshot*. Rather than embed its interactive viewer, we reconstruct the\n",
    "one plot that matters — allocated bytes over the step — directly from the\n",
    "recorded trace, and see the sawtooth of the figure in\n",
    "real data:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-4 fw=all -->\n",
    "\n",
    "Each cycle is one training step: memory climbs through the forward pass as\n",
    "activations accumulate, then falls through the backward pass as they are\n",
    "released — exactly the anatomy, now measured. The JAX side plans instead\n",
    "of counts. Ahead-of-time compilation (that section) hands\n",
    "back an object whose `memory_analysis()` reports what the compiler\n",
    "*decided* to allocate, before any memory is touched:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-measuring-memory-3"
   },
   "source": [
    "def loss_fn(params, X):\n",
    "    h = X\n",
    "    for W in params:\n",
    "        h = jax.nn.relu(h @ W)\n",
    "    return h.sum()\n",
    "\n",
    "key = jax.random.PRNGKey(0)\n",
    "params = [jax.random.normal(k, (2048, 2048)) * 0.02\n",
    "          for k in jax.random.split(key, 6)]\n",
    "X = jax.random.normal(key, (256, 2048))\n",
    "\n",
    "compiled = jax.jit(jax.grad(loss_fn)).lower(params, X).compile()\n",
    "a = compiled.memory_analysis()\n",
    "print(f'compiler-planned temp memory: {a.temp_size_in_bytes / 1e6:.0f} MB')\n",
    "print(f'argument + output: '\n",
    "      f'{(a.argument_size_in_bytes + a.output_size_in_bytes) / 1e6:.0f} MB')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "compiler-planned temp memory: 15 MB\nargument + output: 203 MB\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-5 fw=all -->\n",
    "\n",
    "The contrast is worth stating plainly: **PyTorch counts allocations as\n",
    "they happen; XLA plans memory at compile time.** Neither is better; they\n",
    "reflect the eager-versus-traced split of that section, and\n",
    "knowing which mental model your framework uses tells you where to look\n",
    "when memory surprises you. Just do not compare the two frameworks' digits\n",
    "head-to-head: a compiler's plan is not an allocator's high-water mark.\n",
    "\n",
    "## Mixed Precision\n",
    "\n",
    "The format ladder of that section promised that every halving of\n",
    "width wins twice — double the tensor-core throughput and half the bytes.\n",
    "Mixed precision cashes that promise for training. The idea: run the\n",
    "compute-heavy operations (matmuls, convolutions) in bf16, where tensor\n",
    "cores are fastest and activations are half-size, while keeping a few\n",
    "numerically delicate things (the master weights, the loss reduction) in\n",
    "fp32. On our Ada card the arithmetic win is real and robust — but only if\n",
    "you know what your baseline is running, the lesson that section\n",
    "flagged. Whether plain fp32 matmuls use the tensor cores' tf32 path is a\n",
    "framework default: PyTorch ships with tf32 *off*, so a naive fp32\n",
    "baseline is needlessly slow and inflates the apparent speedup; JAX on\n",
    "this card runs tf32-class dot compute *by default*, so its naive baseline\n",
    "is already fair — but is not the true fp32 it appears to be. We time all\n",
    "three configurations, once, so the difference is unmistakable:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-mixed-precision"
   },
   "source": [
    "def net_fn(params, X, dtype):\n",
    "    h = X.reshape(X.shape[0], -1).astype(dtype)\n",
    "    for W in params[:-1]:\n",
    "        h = jax.nn.relu(h @ W.astype(dtype))\n",
    "    return h @ params[-1].astype(dtype)\n",
    "\n",
    "key = jax.random.PRNGKey(0)\n",
    "sizes = [(64 * 64, 8192), (8192, 8192), (8192, 10)]\n",
    "params = [jax.random.normal(k, s) * 0.02\n",
    "          for k, s in zip(jax.random.split(key, 3), sizes)]\n",
    "X = jax.random.normal(key, (2048, 64 * 64))\n",
    "\n",
    "def loss_fn(params, dtype):\n",
    "    return (net_fn(params, X, dtype) ** 2).sum()\n",
    "\n",
    "grad_fn = jax.jit(jax.grad(loss_fn), static_argnums=1)\n",
    "with jax.default_matmul_precision('highest'):  # Unfair baseline: tf32 off\n",
    "    print(d2l.Benchmark(lambda: grad_fn(params, jnp.float32),\n",
    "                        desc='fp32 (tf32 off)'))\n",
    "# The default on this card is tf32-class compute -- the fair baseline\n",
    "print(d2l.Benchmark(lambda: grad_fn(params, jnp.float32), desc='tf32'))\n",
    "print(d2l.Benchmark(lambda: grad_fn(params, jnp.bfloat16), desc='bf16'))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "fp32 (tf32 off): 26.41 ms/call\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "tf32: 15.99 ms/call\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "bf16: 8.70 ms/call\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-6 fw=all -->\n",
    "\n",
    "Against the fair tf32 baseline, bf16 runs about one and a half to two\n",
    "times as fast here — the exact ratio is framework- and shape-dependent —\n",
    "and the three timings tell the same two-step story in both tabs: moving\n",
    "true fp32 dots onto the tf32 tensor cores already buys a sizable step,\n",
    "and bf16 then adds at least as much again. Both are genuine tensor-core\n",
    "wins, not measurement artifacts.\n",
    "Note what the PyTorch tab does *not* use: a `GradScaler`. Loss scaling\n",
    "exists to keep tiny fp16 gradients from underflowing fp16's narrow\n",
    "5-bit-exponent range; bf16 shares fp32's 8-bit exponent\n",
    "(the figure), so its gradients do not underflow and no\n",
    "scaler is needed. fp16-plus-`GradScaler` is the pre-Ampere legacy path,\n",
    "worth one sentence and an exercise. The JAX tab makes the philosophical\n",
    "difference visible: precision in JAX is *explicit* — you thread dtypes\n",
    "through the computation, and `jax.default_matmul_precision` sets how fp32\n",
    "dots are *computed* (true fp32 versus tf32-class tensor-core compute)\n",
    "without changing what any tensor *stores*. On this card JAX's default\n",
    "matches `'high'`, tf32-class, bit for bit — which is why the middle\n",
    "timing is the fair baseline. PyTorch decides per-operation what to cast;\n",
    "JAX makes you say it. Both land on the same design point: bf16 halves the\n",
    "bytes and doubles the tensor-core rate.\n",
    "\n",
    "## Activation Checkpointing\n",
    "\n",
    "Mixed precision halves activation memory; activation checkpointing can cut\n",
    "it much further, by refusing to *store* most activations at all. Recall\n",
    "from the figure that activations are held from the\n",
    "moment the forward pass produces them until the backward pass consumes\n",
    "them — that storage is the peak. Checkpointing makes a different trade:\n",
    "store only a few activations (at block boundaries), and when the backward\n",
    "pass needs the ones in between, *recompute them* by re-running that block's\n",
    "forward pass. You pay extra compute — at most one extra forward pass — to\n",
    "avoid holding a large fraction of the activations; the cells below\n",
    "measure both sides of the trade. This is the same\n",
    "\"recompute rather than store\" argument the\n",
    "Mamba kernel of that section made one level down in the hardware\n",
    "[@Chen.Xu.Zhang.ea.2016]."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-activation-checkpointing"
   },
   "source": [
    "def block(W1, W2, x):\n",
    "    return x + jax.nn.gelu(x @ W1) @ W2\n",
    "\n",
    "key = jax.random.PRNGKey(1)\n",
    "ks = jax.random.split(key, 32)             # one key per weight matrix\n",
    "Ws = [(jax.random.normal(ks[2 * i], (1024, 1024)) * 0.02,\n",
    "       jax.random.normal(ks[2 * i + 1], (1024, 1024)) * 0.02)\n",
    "      for i in range(16)]\n",
    "X = jax.random.normal(key, (16384, 1024))\n",
    "\n",
    "# jax.checkpoint (a.k.a. remat) recomputes the block in the backward pass.\n",
    "ckpt_block = jax.checkpoint(block)\n",
    "def forward(Ws, X, blk):\n",
    "    h = X\n",
    "    for W1, W2 in Ws:\n",
    "        h = blk(W1, W2, h)\n",
    "    return (h ** 2).sum()\n",
    "\n",
    "for name, blk in [('store all', block), ('checkpointed', ckpt_block)]:\n",
    "    compiled = jax.jit(jax.grad(forward), static_argnums=2).lower(\n",
    "        Ws, X, blk).compile()\n",
    "    mb = compiled.memory_analysis().temp_size_in_bytes / 1e6\n",
    "    t = d2l.Benchmark(lambda: compiled(Ws, X)).time\n",
    "    print(f'{name}: compiler temp {mb:.0f} MB, fwd+bwd {1000 * t:.1f} ms')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "store all: compiler temp 3230 MB, fwd+bwd 50.5 ms\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "checkpointed: compiler temp 1216 MB, fwd+bwd 58.7 ms\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-7 fw=all -->\n",
    "\n",
    "Checkpointing releases about 2 GB of held activations on this 16-block\n",
    "stack — the measured peak falls by nearly half in the allocator's count\n",
    "and by more than half in the compiler's plan — and the completion-timed\n",
    "step runs only about 15–20% longer, comfortably under the one-extra-forward bound, since the\n",
    "backward pass dominates the step. It is a trade you make when memory is\n",
    "the binding constraint: checkpointing normally costs time, and its\n",
    "purpose is memory (the rare configurations where recomputation beats a\n",
    "cache-hostile stash do not change the rule). JAX exposes the same\n",
    "mechanism as `jax.checkpoint` (equivalently `jax.remat`), with a nicer\n",
    "knob: policies like `checkpoint_dots` let you say *which* intermediates to\n",
    "keep, and `jax.remat`'s companion tools let you *watch* what was saved —\n",
    "the introspection theme JAX does uniquely well. The trade is identical;\n",
    "the interface reflects each framework's temperament, imperative versus\n",
    "declarative.\n",
    "\n",
    "## Gradient Accumulation\n",
    "\n",
    "The final technique buys the *statistical* effect of a large batch without\n",
    "its memory cost — and discharges a promise made all the way back in\n",
    "that section. A large batch improves gradient estimates and\n",
    "tensor-core utilization (its intensity is higher, that section),\n",
    "but a large batch's activations may not fit. Gradient accumulation splits\n",
    "the desired *global* batch into $k$ *micro-batches*, runs forward and\n",
    "backward on each while *summing* the gradients, and only then takes one\n",
    "optimizer step. The activations of just one micro-batch are live at a\n",
    "time, so peak memory follows the micro-batch, while the update sees the\n",
    "full global batch: $B_{\\textrm{global}} = B_{\\textrm{micro}} \\times k$\n",
    "(times the number of devices, once that section adds them). The\n",
    "parity check — that accumulating $k$ micro-batches matches one full-batch\n",
    "step — is worth seeing, because it is the whole correctness claim:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-gradient-accumulation"
   },
   "source": [
    "W = jax.random.normal(jax.random.PRNGKey(0), (1024, 1))\n",
    "X = jax.random.normal(jax.random.PRNGKey(1), (128, 1024))\n",
    "y = jax.random.normal(jax.random.PRNGKey(2), (128, 1))\n",
    "\n",
    "def mse(W, X, y):\n",
    "    return ((X @ W - y) ** 2).mean()\n",
    "\n",
    "full = jax.grad(mse)(W, X, y)\n",
    "acc = sum(jax.grad(mse)(W, X[i * 32:(i + 1) * 32], y[i * 32:(i + 1) * 32])\n",
    "          for i in range(4)) / 4\n",
    "print(f'max gradient difference: {jnp.abs(full - acc).max():.2e}')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "max gradient difference: 3.10e-06\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-8 fw=all -->\n",
    "\n",
    "The gradients match to floating-point noise. The one subtlety is the\n",
    "averaging: if your loss already averages over the batch (as MSE and\n",
    "cross-entropy do), each micro-batch gradient must be scaled by $1/k$ so\n",
    "the sum reproduces the mean. The equivalence also rests on conditions\n",
    "this tiny example satisfies silently: an additive per-example loss,\n",
    "equal-size micro-batches (or a properly weighted mean), the same dropout\n",
    "and RNG draws, no batch-coupled state such as BatchNorm — whose\n",
    "statistics would be computed per micro-batch — and clipping or optimizer\n",
    "steps applied only at the global boundary. Accumulation is the right tool\n",
    "when you are memory-bound and want a bigger *effective* batch; it is the\n",
    "wrong tool when you want to go faster: the arithmetic is identical but\n",
    "arrives as more launches over smaller, lower-intensity GEMMs, so by this\n",
    "chapter's own accounting it is usually slightly *slower* — it only\n",
    "rearranges *when* the optimizer step happens.\n",
    "\n",
    "## The Ladder So Far\n",
    "\n",
    "Four sections in, the escalation is worth naming as a checklist. When a\n",
    "model does not fit or does not run fast enough, in order of cost:\n",
    "\n",
    "1. **It's slow, bandwidth- or overhead-bound** → compile it\n",
    "   (that section). Usually free.\n",
    "2. **It's slow, compute-bound, or it's tight on memory** → drop to bf16\n",
    "   (mixed precision). Around $1.5$–$2\\times$ speed and half the\n",
    "   activation bytes.\n",
    "3. **It still doesn't fit** → checkpoint activations. Large memory saving\n",
    "   for a modest increase in time.\n",
    "4. **It still doesn't fit, and you want a bigger effective batch** →\n",
    "   accumulate gradients. Bigger effective batch at micro-batch memory —\n",
    "   about the same wall-clock, never faster.\n",
    "5. **It's still too slow, or genuinely too big for one card** → add\n",
    "   devices (that section).\n",
    "\n",
    "The next section takes that last step — and finds that adding devices is\n",
    "where the accounting gets hardest.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* A training step's memory is parameters ($4P$) + gradients ($4P$) +\n",
    "  optimizer states (Adam: $8P$) + activations (batch × depth × width, the\n",
    "  term that explodes). About $16P$ bytes before activations; peak lands\n",
    "  early in the backward pass.\n",
    "* PyTorch measures memory with runtime counters and allocation snapshots;\n",
    "  JAX reports the compiler's ahead-of-time memory *plan*. Counting versus\n",
    "  planning — the eager-versus-traced split again.\n",
    "* Mixed-precision bf16 raises tensor-core throughput (1.5–2× over a fair\n",
    "  tf32 baseline here, framework- and shape-dependent) and halves\n",
    "  activation bytes; measure it against that *fair* tf32 baseline — the\n",
    "  frameworks' fp32 defaults differ — and skip the `GradScaler` (bf16\n",
    "  shares fp32's exponent range — only fp16 needs loss scaling).\n",
    "* Activation checkpointing recomputes instead of stores: peak memory cut\n",
    "  roughly in half or better on a deep stack (about 2 GB of activations\n",
    "  released) for about 15–20% more step time here. It normally costs time; its\n",
    "  purpose is memory — use it when memory binds.\n",
    "* Gradient accumulation reaches a large *effective* batch at micro-batch\n",
    "  memory cost — $B_{\\textrm{global}} = B_{\\textrm{micro}} \\times k$ — at\n",
    "  the same wall-clock; scale per-micro-batch losses by $1/k$ when the loss\n",
    "  already averages.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Take the GPT of that section and, using only the anatomy of\n",
    "   that section, compute the largest $(\\textrm{width},\n",
    "   \\textrm{depth})$ trainable in 24 GB under three regimes: fp32 + Adam,\n",
    "   bf16 + Adam, bf16 + Adam + checkpointing (assume activations dominate).\n",
    "   Then verify one point against `max_memory_allocated`.\n",
    "1. Derive the per-token activation bytes of one `TransformerBlock`\n",
    "   (that section) as a function of width and sequence length, and\n",
    "   check your formula against the measured peak from the snapshot plot.\n",
    "1. Implement \"checkpoint every $\\sqrt{n}$-th block\" for the deep block\n",
    "   stack above and measure peak memory and step time. Why is $\\sqrt{n}$ the\n",
    "   memory-optimal checkpointing interval?\n",
    "1. Break fp16 (not bf16) deliberately: construct a step whose gradients\n",
    "   are small enough to underflow fp16's subnormal floor (e.g., scale the\n",
    "   loss by $10^{-4}$ before calling backward), run it under\n",
    "   `torch.autocast(dtype=torch.float16)` without a `GradScaler`, and\n",
    "   verify that gradients flush to zero while the same step under bf16\n",
    "   survives. Then enable a `GradScaler`, confirm the gradients return,\n",
    "   and explain what the scaler did in terms of\n",
    "   the figure. Why does a full training run fail this\n",
    "   way only *sometimes* — what makes the failure workload-dependent?\n",
    "1. With a fixed global batch of 512, sweep the micro-batch size $\\in\n",
    "   \\{512, 256, 128, 64\\}$ (so $k \\in \\{1, 2, 4, 8\\}$), and plot peak\n",
    "   memory and wall-clock per optimizer step. Confirm that memory follows\n",
    "   the micro-batch while the update is unchanged, and explain the small\n",
    "   time overhead of large $k$."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
