{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-1 fw=all -->\n",
    "\n",
    "# Reproducibility and Inspection\n",
    "\n",
    "Run yesterday's experiment again and the loss curve comes out different.\n",
    "Is this morning's change an improvement, or a lucky seed? Answering that\n",
    "question requires knowing where every random number in a training run comes\n",
    "from. And once two runs *do* disagree, or a loss turns into NaN, the next\n",
    "question is what the model computed layer by layer, ideally without editing\n",
    "its source. This section covers both skills: controlling randomness (seeds,\n",
    "generators, determinism) and observing a running model from the outside\n",
    "(hooks)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-reproducibility-and-inspection"
   },
   "source": [
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "import optax"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-2 fw=all -->\n",
    "\n",
    "## Seeds and Randomness\n",
    "\n",
    "Randomness enters a training run in more places than you might list on a\n",
    "first try: initialization draws every weight (that section),\n",
    "dropout samples a fresh mask at each step\n",
    "[@Srivastava.Hinton.Krizhevsky.ea.2014], the data loader shuffles\n",
    "examples differently in every epoch, augmentations sample crops and flips,\n",
    "and the loader's worker processes each do all of this in parallel. Every one\n",
    "of these consults a pseudorandom number generator, a deterministic algorithm\n",
    "whose entire output sequence is fixed by its seed. Seed every generator and\n",
    "the run is repeatable; miss one and it is not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-3 fw=jax -->\n",
    "\n",
    "JAX has no global generator to seed. Every random operation takes a *key*,\n",
    "a value that fully determines its output: `jax.random.normal(key, shape)`\n",
    "returns the same numbers every time it is called with that key, and fresh\n",
    "randomness comes only from deriving fresh keys with `jax.random.split`. The\n",
    "function below turns its seed into one key and splits it three ways, one\n",
    "key each for initialization, inputs, and targets, so the seed pins the\n",
    "whole run by construction:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-seeds-and-randomness"
   },
   "source": [
    "def train_once(seed):\n",
    "    key_init, key_X, key_y = jax.random.split(jax.random.key(seed), 3)\n",
    "    rngs = nnx.Rngs(key_init)\n",
    "    net = nnx.Sequential(nnx.Linear(20, 32, rngs=rngs),\n",
    "                         nnx.relu,\n",
    "                         nnx.Linear(32, 1, rngs=rngs))\n",
    "    X = jax.random.normal(key_X, (128, 20))\n",
    "    y = jax.random.normal(key_y, (128, 1))\n",
    "    optimizer = nnx.Optimizer(net, optax.sgd(0.1), wrt=nnx.Param)\n",
    "    loss_fn = lambda model: ((model(X) - y) ** 2).mean()\n",
    "    for _ in range(5):\n",
    "        loss, grads = nnx.value_and_grad(loss_fn)(net)\n",
    "        optimizer.update(net, grads)\n",
    "    return loss, net.layers[0].kernel[...]\n",
    "\n",
    "loss_a, w_a = train_once(seed=0)\n",
    "loss_b, w_b = train_once(seed=0)\n",
    "loss_c, _ = train_once(seed=1)\n",
    "print('same seed, identical loss and weights:',\n",
    "      loss_a == loss_b, jnp.array_equal(w_a, w_b))\n",
    "print('different seed:', loss_a, 'vs', loss_c)"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "same seed, identical loss and weights: True True\ndifferent seed: 1.0159674 vs 0.79100513\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-4 fw=all -->\n",
    "\n",
    "The two seeded runs agree *bitwise*, down to the last floating-point bit:\n",
    "same initial weights, same data, same gradients, same operations in the\n",
    "same order. One detail deserves a flag: we pinned the TensorFlow run to\n",
    "the CPU. On a GPU the same seeded code produces equal losses but weights\n",
    "that differ in their last bits — the reason emerges in the next section.\n",
    "\n",
    "### Generator Objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-5 fw=jax -->\n",
    "\n",
    "A single shared stream is fragile: every consumer advances it, so\n",
    "inserting one extra random call shifts everything drawn after it. In JAX\n",
    "the keys *are* the generator objects: each key is a private stream you\n",
    "hold as a value, and `jax.random.split` manufactures as many independent\n",
    "streams as you need. The property a private stream is supposed to\n",
    "protect, that unrelated draws elsewhere cannot shift yours, holds by\n",
    "construction, because no draw advances any shared state. Here an\n",
    "unrelated draw between two uses of the same key changes nothing:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-generator-objects"
   },
   "source": [
    "key_split, key_other = jax.random.split(jax.random.key(42))\n",
    "split = jax.random.permutation(key_split, 10)\n",
    "_ = jax.random.normal(key_other, (1000,))  # unrelated draw, its own key\n",
    "split_again = jax.random.permutation(key_split, 10)\n",
    "print(jnp.array_equal(split, split_again), split)"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "True [1 3 9 5 0 2 6 7 4 8]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-6 fw=jax -->\n",
    "\n",
    "Every function in `jax.random` takes the key as its first argument; there\n",
    "is no variant that consults hidden state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-7 fw=all -->\n",
    "\n",
    "### DataLoader Workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-8 fw=jax -->\n",
    "\n",
    "Parallel loader workers still need separate random streams. With explicit\n",
    "keys, those streams are visible in the program: a worker's randomness is\n",
    "whatever key you hand it, so you split one key\n",
    "into per-worker keys (`jax.random.split(key, num_workers)`), refresh them\n",
    "each epoch by folding in the epoch number (`jax.random.fold_in`), and two\n",
    "workers can share a stream only if the code visibly passes the same key\n",
    "twice. JAX programs usually borrow their input\n",
    "pipeline from NumPy, PyTorch, or `tf.data`, and those loaders bring their\n",
    "own worker-seeding contract. When you do that, configure the loader's\n",
    "workers from one base seed and refresh them each epoch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-9 fw=all -->\n",
    "\n",
    "### Randomness as a Value\n",
    "\n",
    "Every failure above is hidden global state: a generator that lives\n",
    "somewhere your code does not mention, silently copied or shared. Some\n",
    "library designs remove the hiding place altogether. In an explicit-PRNG\n",
    "design (JAX is the prominent example), randomness is threaded through the\n",
    "program as a value: every random operation takes a key argument, using the\n",
    "same key twice yields the same numbers by definition, and independent\n",
    "streams are obtained by explicitly splitting a key in two. Accidentally\n",
    "duplicating a stream would require writing the duplication into the code,\n",
    "so the worker bug cannot be expressed. The price is bookkeeping, since\n",
    "every function that needs randomness must accept and split keys. PyTorch's\n",
    "global seed is the convenient version of the same contract, one implicit\n",
    "key that everything shares; generator objects recover the explicit version\n",
    "where it matters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-10 fw=jax -->\n",
    "\n",
    "The two cells above are this design in action: `train_once` turned one\n",
    "seed into one key and split it, and reusing a key reproduced the\n",
    "permutation bit for bit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-11 fw=all -->\n",
    "\n",
    "## Determinism and Its Price\n",
    "\n",
    "Seeding fixes which numbers the program draws. It does not fix how the\n",
    "arithmetic evaluates. Floating-point addition is not associative\n",
    "(that section), so summing the same numbers in a different\n",
    "order or grouping may give a different answer:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-determinism-and-its-price-1"
   },
   "source": [
    "x = jax.random.normal(jax.random.key(0), (1_000_000,))\n",
    "s_fwd, s_rev = x.sum(), x[::-1].sum()  # same numbers, different order\n",
    "print(s_fwd == s_rev, (s_fwd - s_rev).item())"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "False 0.000213623046875\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-12 fw=jax -->\n",
    "\n",
    "For the *drawn numbers*, JAX's answer is structural. Its PRNG is\n",
    "counter-based (Threefry): a draw is a pure function of the key and the\n",
    "requested shape, not of any evolving state, so the same key produces the\n",
    "same numbers on CPU, GPU, and TPU alike (the documented caveat is that\n",
    "JAX does not promise identical bits *across its own releases*). For the\n",
    "*arithmetic*, the story splits by backend. On CPU, XLA's kernels are\n",
    "deterministic, and the bitwise agreement of `train_once` above is exactly\n",
    "that, verified. On GPU, some XLA kernels commit partial sums via atomic\n",
    "additions in whatever order threads finish, the same last-bits\n",
    "nondeterminism as elsewhere; setting the environment variable\n",
    "`XLA_FLAGS=--xla_gpu_deterministic_ops=true` before JAX initializes\n",
    "forces deterministic implementations, at a speed cost we have not\n",
    "measured here. There is no error-raising mode: the flag substitutes\n",
    "deterministic kernels rather than refusing nondeterministic ones."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-13 fw=all -->\n",
    "\n",
    "This guarantee has a narrow scope. No framework promises bitwise agreement\n",
    "across releases, platforms, or CPU versus GPU: it holds only for a\n",
    "pinned machine, library version, and flag configuration. That makes bitwise\n",
    "reproducibility a *debugging* tool, the setting that lets you bisect\n",
    "exactly where two runs diverge. The *scientific* goal is statistical\n",
    "reproducibility: the same conclusions across seeds, reported as a mean and\n",
    "spread over several runs rather than one fortunate curve. The distinction\n",
    "mirrors that section: changing dtype changes results in the\n",
    "last bits by design, and an experimental claim that survives neither a new\n",
    "seed nor bfloat16 was never a result.\n",
    "\n",
    "## Hooks: Looking Inside"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-14 fw=jax -->\n",
    "\n",
    "NNX can wrap a module call with `nnx.capture`, recording the return value of\n",
    "each submodule method alongside the model output without changing the model's\n",
    "source. That matters\n",
    "precisely when the model came from a library or a checkpoint you do not\n",
    "want to edit. the figure draws the general picture: an\n",
    "observation point in the gap the call wrapper already leaves around each\n",
    "module's computation, from which an observer can capture or check\n",
    "without a single line of the model changing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-15 fw=all -->\n",
    "\n",
    "![The `__call__` wrapper as a pipeline: input flows through pre-hooks, then forward, then hooks, to the output, with the two hook stages dashed and orange against forward's solid blue, and a side arrow from the hooks stage to an observer that can capture, check, or modify.](../img/bg-hooks.svg)\n",
    "\n",
    "We reuse the residual stack of\n",
    "that section, rebuilt compactly:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-hooks-looking-inside"
   },
   "source": [
    "class ResidualBlock(nnx.Module):\n",
    "    def __init__(self, num_hiddens, rngs):\n",
    "        self.body = nnx.Sequential(\n",
    "            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs), nnx.relu,\n",
    "            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs))\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return X + self.body(X)\n",
    "\n",
    "rngs = nnx.Rngs(0)\n",
    "net = nnx.Sequential(nnx.Linear(20, 64, rngs=rngs),\n",
    "                     *[ResidualBlock(64, rngs) for _ in range(8)],\n",
    "                     nnx.Linear(64, 10, rngs=rngs))\n",
    "X = jax.random.normal(jax.random.key(1), (256, 20))"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-16 fw=all -->\n",
    "\n",
    "### Capturing Activation Statistics\n",
    "\n",
    "The initialization experiments of that section measured the\n",
    "standard deviation of activations at every depth. The same measurement\n",
    "takes a few lines on an unmodified model:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-capturing-activation-statistics"
   },
   "source": [
    "_, inter = nnx.capture(\n",
    "    net, nnx.Intermediate, method_outputs=nnx.Intermediate)(X)\n",
    "for k, layer in enumerate(net.layers):\n",
    "    out = inter['layers'][k]['__call__'][0]\n",
    "    print(f'{type(layer).__name__:15s} std {out.std():.2f}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Linear          std 0.99\nResidualBlock   std 1.23\nResidualBlock   std 1.50\nResidualBlock   std 1.78\nResidualBlock   std 2.15\nResidualBlock   std 2.75\nResidualBlock   std 3.26\nResidualBlock   std 3.90\nResidualBlock   std 4.33\nLinear          std 4.27\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-17 fw=all -->\n",
    "\n",
    "The residual stream's spread grows block by block, since each block adds\n",
    "its body's output on top of the stream, exactly the depth effect that\n",
    "motivated the scaled initializations of that section, measured\n",
    "here without touching the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-18 fw=jax -->\n",
    "\n",
    "Nothing here needs detaching or removing: the captured intermediates are\n",
    "ordinary arrays with no autograd graph attached, the dictionary exists\n",
    "for this one call, and no observer stays registered on the model. The\n",
    "`method_outputs=nnx.Intermediate` argument records every method return. For\n",
    "finer control, a module can call\n",
    "`self.sow(nnx.Intermediate, 'name', value)` at selected points, and\n",
    "`nnx.capture(model, nnx.Intermediate)` then returns only those named values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-19 fw=all -->\n",
    "\n",
    "### A NaN Finder\n",
    "\n",
    "When a loss becomes NaN at step 40,000, the useful question is which layer\n",
    "produced the first non-finite value. Check every leaf module's output for\n",
    "finiteness and the answer arrives in one forward pass. We sabotage a\n",
    "weight deep in the stack to watch it fire:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-a-nan-finder"
   },
   "source": [
    "saved = net.layers[3].body.layers[0].kernel[0, 0]\n",
    "net.layers[3].body.layers[0].kernel[0, 0] = float('nan')\n",
    "_, inter = nnx.capture(\n",
    "    net, nnx.Intermediate, method_outputs=nnx.Intermediate)(X)\n",
    "\n",
    "def get_path(tree, path):\n",
    "    for key in path:\n",
    "        tree = tree[key]\n",
    "    return tree\n",
    "\n",
    "for path, module in nnx.iter_modules(net):\n",
    "    if isinstance(module, nnx.Linear):\n",
    "        out = get_path(inter, path)['__call__'][0]\n",
    "        if not jnp.isfinite(out).all():\n",
    "            print('first non-finite output in', path)\n",
    "            break\n",
    "\n",
    "net.layers[3].body.layers[0].kernel[0, 0] = saved  # net is shared; undo the sabotage"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "first non-finite output in ('layers', 3, 'body', 'layers', 0)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-20 fw=jax -->\n",
    "\n",
    "We inspect the captured outputs of linear modules in object-graph order, which\n",
    "matches execution order for this sequential network. The report names\n",
    "`('layers', 3, 'body', 'layers', 0)`, the layer we poisoned, rather than\n",
    "leaving you to bisect with print statements while NaNs propagate through\n",
    "everything downstream."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-21 fw=all -->\n",
    "\n",
    "### Backward Hooks and Beyond"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-22 fw=jax -->\n",
    "\n",
    "Gradients need no hook at all, because they are not events that fire\n",
    "inside a backward pass: `nnx.grad` returns them as a tree of values in\n",
    "the same shape as the parameters. To log per-layer gradient norms, catch\n",
    "exploding gradients at their source, or experiment with per-layer\n",
    "clipping, compute the gradient tree and inspect or transform it like any\n",
    "other data:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "reproducibility-inspection-backward-hooks-and-beyond"
   },
   "source": [
    "grads = nnx.grad(lambda model: (model(X) ** 2).mean())(net)\n",
    "norms = jax.tree_util.tree_map(jnp.linalg.norm, grads)\n",
    "print(norms['layers'][3]['body']['layers'][0])"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n  \u001b[38;2;156;220;254m'bias'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(1.9621452, dtype=float32)\n  \u001b[38;2;255;213;3m)\u001b[0m,\n  \u001b[38;2;156;220;254m'kernel'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(15.665125, dtype=float32)\n  \u001b[38;2;255;213;3m)\u001b[0m\n\u001b[38;2;255;213;3m})\u001b[0m\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-23 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "Reproducibility is an inventory problem: initialization, dropout,\n",
    "shuffling, augmentation, and loader workers each draw from some generator,\n",
    "and a run repeats only if every one of them is seeded."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-24 fw=jax -->\n",
    "\n",
    "In JAX the inventory collapses to one item, the key you pass: the same\n",
    "key gives the same draws by construction of the counter-based PRNG, and\n",
    "independent streams come from `jax.random.split`, never from hidden\n",
    "state. Keys make the program repeatable, not the arithmetic: on CPU,\n",
    "XLA's kernels are already deterministic, while on GPU the flag\n",
    "`--xla_gpu_deterministic_ops=true` pins kernel choice too."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-25 fw=all -->\n",
    "\n",
    "Bitwise agreement is a debugging tool;\n",
    "conclusions that hold across seeds are the scientific goal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-26 fw=jax -->\n",
    "\n",
    "For looking inside a model, `nnx.capture(..., method_outputs=...)` returns\n",
    "submodule outputs from an unmodified call, `sow` records named values from\n",
    "the inside, and gradients are values from `nnx.grad` you inspect\n",
    "directly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-27 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Extend `train_once` to load data through your framework's parallel data\n",
    "   loader (e.g., PyTorch's `DataLoader`, with `num_workers=4`). Give the\n",
    "   dataset its own `np.random.default_rng()` object and use it to add noise\n",
    "   as each example loads. On a process-based loader, check whether workers\n",
    "   copied the same generator state. Using the per-worker seeding hook this\n",
    "   section introduced for your framework, replace the dataset's stashed\n",
    "   generator with one derived from the worker's own seed, then verify\n",
    "   agreement across two runs.\n",
    "2. Write a forward hook (or your framework's equivalent introduced in this\n",
    "   section) that counts multiply-accumulate operations for every linear\n",
    "   layer from the shapes of its input and weight, and use it to report\n",
    "   per-layer and total FLOPs for the residual stack above. Check the total\n",
    "   against a hand count.\n",
    "3. Using a backward hook (or your framework's equivalent introduced in this\n",
    "   section), record the norm of each block's output gradient. After the\n",
    "   backward pass, compare these activation-gradient norms with the\n",
    "   parameter-gradient norms in the same block. Which one first reveals an\n",
    "   exploding backward signal?\n",
    "4. A forward hook that returns a value *replaces* the module's output. Use\n",
    "   one to zero out the output of a single residual block's body (turning\n",
    "   the block into the identity) and measure how the network's output\n",
    "   changes. Which block matters most, and how would you find out with one\n",
    "   loop?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=reproducibility-inspection-md-28 fw=jax -->\n",
    "\n",
    "5. Rebuild the activation-statistics table by editing `ResidualBlock` to call\n",
    "   `self.sow(nnx.Intermediate, 'body_out', ...)` on its body's output, then\n",
    "   collect those values with `nnx.capture(net, nnx.Intermediate)`. Compare\n",
    "   capture-all-methods, opt-in `sow`, and\n",
    "   PyTorch-style mutable hooks: which requires touching model\n",
    "   code, which can silently retain memory, and which would you want for a\n",
    "   model you do not own?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
