5.8 Reproducibility and Inspection

Repeated training runs can produce different loss curves because their random streams or arithmetic differ. Evaluating a code change therefore requires an inventory of random-number sources and determinism controls. When two runs disagree or a loss becomes NaN, layer-level inspection can locate the first divergent or non-finite value without changing the model’s computation. This chapter treats reproducibility first, then model inspection through call wrappers and hooks.

import random
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

import tensorflow as tf

import jax
from jax import numpy as jnp
from flax import nnx
import optax

from mxnet import autograd, gluon, init, np, npx
from mxnet.gluon import nn
npx.set_np()

5.8.1 Seeds and Randomness

Randomness enters a training run in more places than you might list on a first try: initialization draws every weight (Section 5.3), dropout samples a fresh mask at each step (Srivastava et al. 2014), the data loader shuffles examples differently in every epoch, augmentations sample crops and flips, and loader workers can perform these operations in parallel. Each random operation consults a pseudorandom number generator whose output sequence is fixed by its seed. Seeding every relevant generator fixes those random streams; bitwise repetition also requires deterministic arithmetic and a stable software and hardware configuration.

In PyTorch, torch.manual_seed(k) seeds the global generator on the CPU and on every CUDA device. It does not touch Python’s random module or NumPy, which keep their own global generators (random.seed, np.random.seed), nor NumPy Generator objects created by np.random.default_rng, which carry their own seeds. The function below draws everything, weights, data, and shuffling, from torch’s global generator, so one call fixes the random draws in this function:

TensorFlow collapses the inventory to one call: tf.keras.utils.set_random_seed(k) seeds Python’s random module, NumPy’s global generator, and TensorFlow’s global generator together, the three a typical training script consults. It does not reach generators that carry their own state (tf.random.Generator objects, NumPy Generator objects from np.random.default_rng), each of which has its own seed. The function below draws everything, weights and data alike, from the seeded global state, so one call fixes the random draws in this function:

JAX has no global generator to seed. Every random operation takes a key, a value that fully determines its output: jax.random.normal(key, shape) returns the same numbers every time it is called with that key, and fresh randomness comes only from deriving fresh keys with jax.random.split. The function below turns its seed into one key and splits it three ways, one key each for initialization, inputs, and targets, so the seed fixes the random draws in this function:

MXNet keeps one global generator per device. np.random.seed(k) (MXNet’s np, an alias of mx.random.seed) seeds all of them in one call, each from the seed and its device id, so initialization, sampling, and dropout are pinned on every device at once (pass device= instead to bring a single generator to a device-independent state). It does not touch Python’s random module or NumPy’s legacy global generator, which keep their own state, a gap that will matter when the data pipeline enters. The function below draws everything, weights and data alike, from MXNet’s seeded generators, so one call fixes the random draws in this function:

def train_once(seed):
    torch.manual_seed(seed)  # seeds the CPU and every CUDA device
    net = nn.Sequential(nn.Linear(20, 32), nn.ReLU(), nn.Linear(32, 1))
    opt = torch.optim.SGD(net.parameters(), lr=0.1)
    X, y = torch.randn(128, 20), torch.randn(128, 1)
    for _ in range(5):
        opt.zero_grad()
        loss = ((net(X) - y) ** 2).mean()
        loss.backward()
        opt.step()
    return loss.item(), net[0].weight.detach().clone()

loss_a, w_a = train_once(seed=0)
loss_b, w_b = train_once(seed=0)
loss_c, _ = train_once(seed=1)
print('same seed, identical loss and weights:',
      loss_a == loss_b, torch.equal(w_a, w_b))
print('different seed:', loss_a, 'vs', loss_c)

same seed, identical loss and weights: True True
different seed: 0.8643209338188171 vs 0.9312677383422852

def train_once(seed):
    tf.keras.utils.set_random_seed(seed)  # seeds Python, NumPy, and TF
    with tf.device('/CPU:0'):  # GPU kernels are not bitwise-deterministic
        net = tf.keras.Sequential([
            tf.keras.layers.Dense(32, activation='relu'),
            tf.keras.layers.Dense(1)])
        opt = tf.keras.optimizers.SGD(learning_rate=0.1)
        X, y = tf.random.normal((128, 20)), tf.random.normal((128, 1))
        for _ in range(5):
            with tf.GradientTape() as tape:
                loss = tf.reduce_mean((net(X) - y) ** 2)
            grads = tape.gradient(loss, net.trainable_variables)
            opt.apply_gradients(zip(grads, net.trainable_variables))
        return float(loss), tf.identity(net.layers[0].kernel)

loss_a, w_a = train_once(seed=0)
loss_b, w_b = train_once(seed=0)
loss_c, _ = train_once(seed=1)
print('same seed, identical loss and weights:',
      loss_a == loss_b, bool(tf.reduce_all(w_a == w_b)))
print('different seed:', loss_a, 'vs', loss_c)

same seed, identical loss and weights: True True
different seed: 1.0948041677474976 vs 1.0242893695831299

def train_once(seed):
    key_init, key_X, key_y = jax.random.split(jax.random.key(seed), 3)
    rngs = nnx.Rngs(key_init)
    net = nnx.Sequential(nnx.Linear(20, 32, rngs=rngs),
                         nnx.relu,
                         nnx.Linear(32, 1, rngs=rngs))
    X = jax.random.normal(key_X, (128, 20))
    y = jax.random.normal(key_y, (128, 1))
    optimizer = nnx.Optimizer(net, optax.sgd(0.1), wrt=nnx.Param)
    loss_fn = lambda model: ((model(X) - y) ** 2).mean()
    for _ in range(5):
        loss, grads = nnx.value_and_grad(loss_fn)(net)
        optimizer.update(net, grads)
    return loss, net.layers[0].kernel[...]

loss_a, w_a = train_once(seed=0)
loss_b, w_b = train_once(seed=0)
loss_c, _ = train_once(seed=1)
print('same seed, identical loss and weights:',
      loss_a == loss_b, jnp.array_equal(w_a, w_b))
print('different seed:', loss_a, 'vs', loss_c)

same seed, identical loss and weights: True True
different seed: 1.0159674 vs 0.79100513

def train_once(seed):
    np.random.seed(seed)  # seeds the generator on every device
    net = nn.Sequential()
    net.add(nn.Dense(32, activation='relu'), nn.Dense(1))
    net.initialize()
    trainer = gluon.Trainer(net.collect_params(),
                            'sgd', {'learning_rate': 0.1})
    X = np.random.normal(0, 1, (128, 20))
    y = np.random.normal(0, 1, (128, 1))
    for _ in range(5):
        with autograd.record():
            loss = ((net(X) - y) ** 2).mean()
        loss.backward()
        trainer.step(1)  # loss is a mean already: no batch-size rescale
    return loss.item(), net[0].weight.data().copy()

loss_a, w_a = train_once(seed=0)
loss_b, w_b = train_once(seed=0)
loss_c, _ = train_once(seed=1)
print('same seed, identical loss and weights:',
      loss_a == loss_b, bool((w_a == w_b).all()))
print('different seed:', loss_a, 'vs', loss_c)

same seed, identical loss and weights: True True
different seed: 1.2009189128875732 vs 0.9397940039634705

The two seeded runs agree bitwise, down to the last floating-point bit: same initial weights, same data, same gradients, same operations in the same order. One detail deserves a flag: we pinned the TensorFlow run to the CPU. On a GPU the same seeded code produces equal losses but weights that differ in their last bits — the reason emerges in the next section.

5.8.1.1 Generator Objects

A single global stream couples every consumer, so inserting one extra random call (a new layer’s initialization, a stray torch.randn while debugging) shifts everything drawn after it. A torch.Generator is a private stream with its own seed. Here a data split stays fixed no matter what else consumes randomness in between:

A single global stream couples every consumer, so inserting one extra random call shifts everything drawn after it. A tf.random.Generator is a private stream with its own seed. Here a data split stays fixed no matter what else consumes randomness in between (the generator has no permutation method, so the example sorts random uniforms):

A single shared stream couples every consumer, so inserting one extra random call shifts everything drawn after it. In JAX the keys are the generator objects: each key is a private stream you hold as a value, and jax.random.split manufactures as many independent streams as you need. The property a private stream is supposed to protect, that unrelated draws elsewhere cannot shift yours, holds by construction, because no draw advances any shared state. Here an unrelated draw between two uses of the same key changes nothing:

A single global stream couples every consumer, so inserting one extra random call shifts everything drawn after it. MXNet offers no private stream to retreat to: there is no generator object, and no sampling function takes one. The substitute is seeding discipline: re-seed immediately before the draws that must be pinned. It works, at a price the private stream does not charge, since re-seeding also resets the stream for everything drawn after it:

g = torch.Generator().manual_seed(42)
split = torch.randperm(10, generator=g)
_ = torch.randn(1000)  # unrelated consumption of the global stream
split_again = torch.randperm(10, generator=g.manual_seed(42))
print(torch.equal(split, split_again), split)

True tensor([2, 6, 1, 8, 4, 5, 0, 9, 3, 7])

g = tf.random.Generator.from_seed(42)
split = tf.argsort(g.uniform((10,)))  # a random permutation from g's stream
_ = tf.random.normal((1000,))  # unrelated consumption of the global stream
split_again = tf.argsort(tf.random.Generator.from_seed(42).uniform((10,)))
print(bool(tf.reduce_all(split == split_again)), split.numpy())

True [4 2 7 3 5 8 9 1 0 6]

key_split, key_other = jax.random.split(jax.random.key(42))
split = jax.random.permutation(key_split, 10)
_ = jax.random.normal(key_other, (1000,))  # unrelated draw, its own key
split_again = jax.random.permutation(key_split, 10)
print(jnp.array_equal(split, split_again), split)

True [1 3 9 5 0 2 6 7 4 8]

np.random.seed(42)
split = np.arange(10)
np.random.shuffle(split)  # mx.np has no permutation; shuffle in place
_ = np.random.normal(0, 1, (1000,))  # unrelated consumption of the stream
np.random.seed(42)  # re-seed at the boundary that must be pinned
split_again = np.arange(10)
np.random.shuffle(split_again)
print(bool((split == split_again).all()), split)

True [9. 1. 4. 5. 8. 7. 6. 2. 3. 0.]

Most sampling functions accept generator=, and so does DataLoader, where it controls the shuffle order. We use that below.

Sampling methods live on the generator itself (g.normal, g.uniform), so a consumer that should own its randomness takes a Generator argument. The input pipeline instead takes a seed directly, which we use next.

Every function in jax.random takes the key as its first argument; there is no variant that consults hidden state.

There is no generator= argument anywhere in Gluon: a consumer that must own its randomness owns a seed instead and calls np.random.seed itself, accepting the reset of everything drawn after it.

5.8.1.2 DataLoader Workers

DataLoader creates a distinct torch seed for every worker from a base seed and the worker id. Current PyTorch also initializes Python’s random module and NumPy’s legacy global generator in each worker. This covers common dataset code, but it cannot discover generator objects owned by a dataset or a third-party library. Those sources still need an explicit worker initializer. The loader’s generator controls both its base seed and the shuffled index order, so one generator can pin the whole input pipeline.

Parallel loader workers still need separate random streams. With explicit keys, those streams are visible in the program: a worker’s randomness is whatever key you hand it, so you split one key into per-worker keys (jax.random.split(key, num_workers)), refresh them each epoch by folding in the epoch number (jax.random.fold_in), and two workers can share a stream only if the code visibly passes the same key twice. JAX programs usually borrow their input pipeline from NumPy, PyTorch, or tf.data, and those loaders bring their own worker-seeding contract. When you do that, configure the loader’s workers from one base seed and refresh them each epoch.

PyTorch

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32  # this worker's torch seed
    np.random.seed(worker_seed)
    random.seed(worker_seed)

g = torch.Generator().manual_seed(0)
data = TensorDataset(torch.arange(8))
# num_workers=0 keeps this notebook cell portable. In a training script,
# num_workers>0 invokes seed_worker once in each worker process.
loader = DataLoader(data, batch_size=4, shuffle=True, num_workers=0,
                    worker_init_fn=seed_worker, generator=g)
print([batch[0].tolist() for batch in loader])
g.manual_seed(0)
print([batch[0].tolist() for batch in loader])

[[5, 3, 1, 7], [2, 6, 4, 0]]
[[5, 3, 1, 7], [2, 6, 4, 0]]

Resetting the generator reproduces the exact shuffle order. With workers enabled, seed_worker derives the legacy NumPy and Python streams from the same base seed. A dataset-owned Generator needs to be replaced there as well.

A reproducibility failure in process-based data loaders occurs when parallel workers inherit or independently reseed a generator: on fork-based loaders every child starts with a byte-for-byte copy of the parent’s NumPy state, so all workers apply the same “random” augmentations, or each child seeds itself from entropy and no two runs agree. tf.data sidesteps the bug class by construction: the pipeline parallelizes with threads inside one process, so there is no forked child to inherit a stale copy, and its randomness is an explicit argument, Dataset.shuffle(buffer, seed=...), not an ambient global that a seeding call may or may not reach. What remains is choosing what the seed means across epochs. With the default reshuffle_each_iteration=True, a seeded pipeline produces a different order on each pass but the same sequence of orders every time the pipeline is rebuilt, fresh shuffles per epoch, repeatable across runs:

TensorFlow

ds = tf.data.Dataset.range(8).shuffle(8, seed=0).batch(4)
print([b.numpy().tolist() for b in ds])  # epoch 1
print([b.numpy().tolist() for b in ds])  # epoch 2: reshuffled, still seeded
ds = tf.data.Dataset.range(8).shuffle(8, seed=0).batch(4)
print([b.numpy().tolist() for b in ds])  # rebuilt pipeline: epoch 1 again

[[5, 6, 7, 2], [1, 3, 4, 0]]
[[4, 6, 7, 2], [0, 5, 3, 1]]
[[5, 6, 7, 2], [1, 3, 4, 0]]

Both alternatives are explicit choices rather than accidents: reshuffle_each_iteration=False freezes one order for every epoch, and leaving seed unset draws fresh orders each run. One knob remains that trades reproducibility away on purpose: parallel map and interleave accept deterministic=False, which hands elements on in completion order for speed; the default True keeps the pipeline’s output order fixed.

Gluon’s DataLoader runs num_workers > 0 on a pool of worker processes, so the fork-inheritance lesson applies verbatim: on Linux each worker starts as a copy of the parent, NumPy state included, and the pool’s initializer does no seeding at all; it only ships the dataset and the set_np flags to the child (see _worker_initializer in gluon/data/dataloader.py). There is no worker_init_fn equivalent to bridge the gap, and the pool outlives any one epoch. Two consequences. First, deliberate per-worker seeding has to live in the dataset itself: make augmentation a deterministic function of the sample by seeding from the sample index inside __getitem__, and it stops mattering which worker runs it. Second, the source shows that the shuffle order of DataLoader(shuffle=True) is drawn from NumPy’s legacy global generator (RandomSampler calls numpy.random.shuffle), so pinning it takes a NumPy seed; MXNet’s own seed does not reach it.

5.8.1.3 Randomness as a Value

The worker failures above arise from generator state that is copied or shared without appearing in the data flow. In an explicit-PRNG design such as JAX, randomness is instead passed through the program as a value: every random operation takes a key, using the same key twice yields the same numbers, and independent streams are obtained by splitting a key. A key can still be reused incorrectly, but the duplication is visible in the program’s data flow. This design requires every function that uses randomness to accept and split keys. PyTorch’s global seed is the convenient version of the same contract, one implicit key that everything shares; generator objects recover the explicit version where it matters.

TensorFlow ships both designs side by side: the stateful API above (a global generator plus tf.random.Generator objects) and a stateless family in which the key is an argument: tf.random.stateless_normal(shape, seed=[k1, k2]) returns the same numbers for the same seed pair by definition. The explicit seed= arguments of tf.data are the same idea applied to the input pipeline.

The two cells above are this design in action: train_once turned one seed into one key and split it, and reusing a key reproduced the permutation bit for bit.

MXNet sits at the fully implicit end of this spectrum: per-device global streams, no generator objects, no stateless variants. Seeding discipline, one call up front and a re-seed at any boundary that must be pinned, is the entire toolkit.

5.8.2 Deterministic Arithmetic and Performance

Seeding fixes which numbers the program draws. It does not fix how the arithmetic evaluates. Floating-point addition is not associative (Section 5.5), so summing the same numbers in a different order or grouping may give a different answer:

torch.manual_seed(0)
x = torch.randn(1_000_000)
s_fwd, s_rev = x.sum(), x.flip(0).sum()  # same numbers, different order
print(s_fwd == s_rev, (s_fwd - s_rev).item())

tensor(False) -0.000244140625

x = tf.random.Generator.from_seed(0).normal((1_000_000,))
s_fwd = tf.reduce_sum(x)
s_rev = tf.reduce_sum(x[::-1])  # same numbers, different order
print(bool(s_fwd == s_rev), float(s_fwd - s_rev))

False -0.000152587890625

x = jax.random.normal(jax.random.key(0), (1_000_000,))
s_fwd, s_rev = x.sum(), x[::-1].sum()  # same numbers, different order
print(s_fwd == s_rev, (s_fwd - s_rev).item())

False -0.0001068115234375

a, b, c = [np.array([v], dtype='float32') for v in (1e8, -1e8, 1.0)]
s_left, s_right = (a + b) + c, a + (b + c)  # same numbers, regrouped
print(bool(s_left == s_right), float(s_left - s_right))

False 1.0

On a CPU the summation order is at least fixed, so seeded runs repeat. On a GPU it often is not: kernels built on atomic additions (scatter_add, some convolution and embedding backward passes) commit their partial sums in whatever order threads happen to finish, so two seeded runs on the same machine can differ in the last bits, and after enough training steps, in the loss curve. torch.use_deterministic_algorithms(True) is the switch that forbids this: operations with a deterministic implementation use it (often at some speed cost), and operations without one raise a RuntimeError rather than returning nondeterministic results. On CUDA you must additionally set the environment variable CUBLAS_WORKSPACE_CONFIG=:4096:8, or cuBLAS matrix multiplications themselves raise. A related but narrower setting is torch.backends.cudnn.benchmark: when true, cuDNN times several convolution algorithms on the first batch and keeps the winner, and since the winner can change from run to run, reproducibility work sets it to False (its sibling cudnn.deterministic covers convolutions only and is subsumed by the global switch).

On a CPU the summation order inside one operation is at least fixed, so seeded runs repeat; the bitwise agreement of train_once above is exactly that. On a GPU it often is not: kernels built on atomic additions (segment sums, the scatter-add behind tf.gather’s gradient) commit their partial sums in whatever order threads happen to finish, so two seeded runs on the same machine can differ in the last bits, and after enough training steps, in the loss curve. tf.config.experimental.enable_op_determinism() is the switch that forbids this: operations with a deterministic implementation use it (often at some speed cost), operations without one raise tf.errors.UnimplementedError rather than returning nondeterministic results, and stateful random operations refuse to run without a seed, since an operation that seeds itself from entropy is nondeterminism by another name. Two properties follow from its design. It is meant to be called at program start, before any operation runs, because it configures kernels as they are created and nothing already executed is redone. And there is no call that turns it off short of a fresh process. We can still demonstrate the seed rule mid-flight by clearing the global seed:

For the drawn numbers, JAX’s answer is structural. Its PRNG is counter-based (Threefry): a draw is a pure function of the key and the requested shape, not of any evolving state, so the same key produces the same numbers on CPU, GPU, and TPU alike (the documented caveat is that JAX does not promise identical bits across its own releases). Arithmetic determinism depends on the backend and operation. The CPU operations used by train_once above produced bitwise agreement in this environment. On GPU, some XLA kernels commit partial sums via atomic additions in whatever order threads finish, the same last-bits nondeterminism as elsewhere; setting the environment variable XLA_FLAGS=--xla_gpu_deterministic_ops=true before JAX initializes forces deterministic implementations, at a speed cost we have not measured here. There is no error-raising mode: the flag substitutes deterministic kernels rather than refusing nondeterministic ones.

We regroup three numbers rather than reversing a million, because MXNet’s sum accumulates in double precision and its result is impressively insensitive to element order — the associativity failure is a property of the format, and the three-line version exhibits it in any framework. On a CPU the summation order inside one operation is at least fixed, so seeded runs repeat; the bitwise agreement of train_once above is exactly that. On a GPU it often is not: kernels built on atomic additions commit their partial sums in whatever order threads happen to finish, so two seeded runs on the same machine can differ in the last bits, and after enough training steps, in the loss curve. MXNet ships no switch that forbids this and no error-raising mode. The one lever it does have is the analogue of cuDNN benchmark mode: by default cuDNN times several convolution algorithms on the first batch and keeps the winner, and since the winner can change from run to run, reproducibility work sets the environment variable MXNET_CUDNN_AUTOTUNE_DEFAULT=0 (before any operator runs) to pin the algorithm choice. Beyond that, staying away from operations whose kernels are nondeterministic is up to you.

torch.use_deterministic_algorithms(True)
if torch.cuda.is_available():
    try:
        x = torch.randn(1, 1, 2, 2, 2, requires_grad=True, device='cuda')
        nn.AvgPool3d(1)(x).sum().backward()
    except RuntimeError as e:
        print(str(e).split('.')[0])
else:
    print('CPU run: every kernel used above is already deterministic;',
          'on CUDA, ops lacking a deterministic kernel raise RuntimeError')
torch.use_deterministic_algorithms(False)

avg_pool3d_backward_cuda does not have a deterministic implementation, but you set 'torch

tf.config.experimental.enable_op_determinism()
tf.random.set_seed(None)  # forget the global seed for a moment
try:
    tf.random.normal((2,))
except RuntimeError as e:
    print(str(e).split('.')[0])
tf.random.set_seed(0)  # determinism stays on; reseed and continue

Random ops require a seed to be set when determinism is enabled

This guarantee has a narrow scope. No framework promises bitwise agreement across releases, platforms, or CPU versus GPU: it holds only for a pinned machine, library version, and flag configuration. That makes bitwise reproducibility a debugging tool, the setting that lets you bisect exactly where two runs diverge. The scientific goal is statistical reproducibility: the same conclusions across seeds, reported as a mean and spread over several runs rather than a single run. The distinction mirrors Section 5.5: changing dtype changes results in the last bits by design, and a conclusion that does not persist across reasonable seeds or numeric configurations requires further evidence.

5.8.3 Inspecting Model Execution with Hooks

Reproducibility determines whether a run can be repeated; inspection determines where a particular run produced an unexpected value. Hooks address the second question without changing the model’s mathematical computation.

In Section 5.1 we noted that net(X) does not call forward directly: it calls __call__, which wraps forward with extra machinery. Hooks are that machinery, exposed. Calling module.register_forward_hook(f) arranges for f(module, inputs, output) to run after every forward pass of that module, with no change to the model’s source, which matters precisely when the model came from a library or a checkpoint you do not want to edit. Figure 5.8.1 draws the wrapper as a pipeline: hooks slot into the gap the __call__ wrapper already leaves around forward, so an observer can capture, check, or modify without a single line of the model changing.

Keras wraps call with a __call__ too (it handles building, dtype casting, and masks), but the wrapper exposes no observation point: nothing can be attached to an unmodified model after the fact, and observing a black-box model from a library or a checkpoint without touching its code has no TensorFlow equivalent. Two idioms reach the same measurements with a little structure. In a functional model every intermediate tensor is a first-class object, so a second tf.keras.Model over the same graph can declare any internal tensor as an output. This secondary model shares all weights and adds no computation. When you own the model’s code, an overridden call can stash or check the requested values as it runs. Figure 5.8.1 still describes the computation order, but TensorFlow requires observation outputs to be declared in the model structure.

NNX can wrap a module call with nnx.capture, recording the return value of each submodule method alongside the model output without changing the model’s source. That matters precisely when the model came from a library or a checkpoint you do not want to edit. Figure 5.8.1 draws the general picture: an observation point in the gap the call wrapper already leaves around each module’s computation, from which an observer can capture or check without a single line of the model changing.

In Section 5.1 we noted that net(X) does not call forward directly: it calls __call__, which wraps forward with extra machinery. Hooks are that machinery, exposed. Calling block.register_forward_hook(f) arranges for f(block, inputs, output) to run after every forward pass of that block, and register_forward_pre_hook attaches before it, with no change to the model’s source, which matters precisely when the model came from a library or a checkpoint you do not want to edit. Gluon itself relies on the mechanism: Block.summary() builds its per-layer table by registering a forward hook on every block. Figure 5.8.1 draws the wrapper as a pipeline: the __call__ wrapper invokes hooks around forward, allowing an observer to capture or check values without modifying the model source (Gluon’s contract is observe-only: a hook’s return value is ignored, so unlike PyTorch’s it cannot modify the output).

Figure 5.8.1: 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.

We reuse the residual stack of Section 5.1, rebuilt compactly:

class ResidualBlock(nn.Module):
    def __init__(self, num_hiddens):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(num_hiddens, num_hiddens),
                                  nn.ReLU(),
                                  nn.Linear(num_hiddens, num_hiddens))

    def forward(self, X):
        return X + self.body(X)

torch.manual_seed(0)
net = nn.Sequential(nn.Linear(20, 64),
                    *[ResidualBlock(64) for _ in range(8)],
                    nn.Linear(64, 10))

def residual_block(X, num_hiddens):
    body = tf.keras.Sequential([
        tf.keras.layers.Dense(num_hiddens, activation='relu'),
        tf.keras.layers.Dense(num_hiddens)])
    return X + body(X)

tf.keras.utils.set_random_seed(0)
inputs = tf.keras.Input(shape=(20,))
taps = [tf.keras.layers.Dense(64)(inputs)]  # keep every intermediate tensor
for _ in range(8):
    taps.append(residual_block(taps[-1], 64))
outputs = tf.keras.layers.Dense(10)(taps[-1])
net = tf.keras.Model(inputs, outputs)

class ResidualBlock(nnx.Module):
    def __init__(self, num_hiddens, rngs):
        self.body = nnx.Sequential(
            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs), nnx.relu,
            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs))

    def __call__(self, X):
        return X + self.body(X)

rngs = nnx.Rngs(0)
net = nnx.Sequential(nnx.Linear(20, 64, rngs=rngs),
                     *[ResidualBlock(64, rngs) for _ in range(8)],
                     nnx.Linear(64, 10, rngs=rngs))
X = jax.random.normal(jax.random.key(1), (256, 20))

class ResidualBlock(nn.Block):
    def __init__(self, num_hiddens):
        super().__init__()
        self.body = nn.Sequential()
        self.body.add(nn.Dense(num_hiddens, activation='relu'),
                      nn.Dense(num_hiddens))

    def forward(self, X):
        return X + self.body(X)

np.random.seed(0)
net = nn.Sequential()
net.add(nn.Dense(64), *[ResidualBlock(64) for _ in range(8)], nn.Dense(10))
net.initialize(init.Xavier())  # variance-preserving, as in the init chapter

5.8.3.1 Capturing Activation Statistics

The initialization experiments of Section 5.3 measured the standard deviation of activations at every depth. The same measurement takes a few lines on an unmodified model:

stats, handles = [], []

def record(module, inputs, output):
    stats.append((type(module).__name__, output.detach().std().item()))

for m in net:
    handles.append(m.register_forward_hook(record))
net(torch.randn(256, 20))
for h in handles:
    h.remove()
for name, std in stats:
    print(f'{name:15s} std {std:.2f}')

Linear          std 0.60
ResidualBlock   std 0.62
ResidualBlock   std 0.66
ResidualBlock   std 0.68
ResidualBlock   std 0.69
ResidualBlock   std 0.72
ResidualBlock   std 0.74
ResidualBlock   std 0.77
ResidualBlock   std 0.78
Linear          std 0.48

probe = tf.keras.Model(inputs, taps + [outputs])  # same layers, same weights
X = tf.random.normal((256, 20))
names = ['Dense'] + ['ResidualBlock'] * 8 + ['Dense']
for name, out in zip(names, probe(X)):
    print(f'{name:15s} std {float(tf.math.reduce_std(out)):.2f}')

Dense           std 0.67
ResidualBlock   std 0.86
ResidualBlock   std 1.03
ResidualBlock   std 1.26
ResidualBlock   std 1.53
ResidualBlock   std 1.86
ResidualBlock   std 2.23
ResidualBlock   std 2.80
ResidualBlock   std 3.60
Dense           std 3.93

_, inter = nnx.capture(
    net, nnx.Intermediate, method_outputs=nnx.Intermediate)(X)
for k, layer in enumerate(net.layers):
    out = inter['layers'][k]['__call__'][0]
    print(f'{type(layer).__name__:15s} std {out.std():.2f}')

Linear          std 0.99
ResidualBlock   std 1.23
ResidualBlock   std 1.50
ResidualBlock   std 1.78
ResidualBlock   std 2.15
ResidualBlock   std 2.75
ResidualBlock   std 3.26
ResidualBlock   std 3.90
ResidualBlock   std 4.33
Linear          std 4.27

stats, handles = [], []

def record(block, inputs, output):
    stats.append((type(block).__name__, output.std().item()))

for m in net:
    handles.append(m.register_forward_hook(record))
net(np.random.normal(0, 1, (256, 20)))
for h in handles:
    h.detach()
for name, std in stats:
    print(f'{name:15s} std {std:.2f}')

Dense           std 0.68
ResidualBlock   std 0.83
ResidualBlock   std 1.03
ResidualBlock   std 1.21
ResidualBlock   std 1.46
ResidualBlock   std 1.83
ResidualBlock   std 2.28
ResidualBlock   std 2.75
ResidualBlock   std 3.43
Dense           std 4.41

The residual stream’s spread grows block by block, since each block adds its body’s output on top of the stream. That is exactly the depth effect that motivated the scaled initializations of Section 5.3, measured here without touching the model.

Two rules keep hooks safe. First, detach before you stash: the hook above stores output.detach().std(); storing output itself would keep the autograd graph of every forward pass alive until stats is cleared, a memory leak that grows with each batch. Second, keep the handle and call handle.remove(): a hook stays registered for the module’s lifetime otherwise, adding overhead to later forward passes and retaining any tensors that it stores.

Nothing needs detaching and nothing needs removing: probe shares the original model’s layers and weights outright, its outputs are ordinary tensors with no gradient tape attached, and no observer stays registered anywhere, because the “hook” is another model output. The limitation is structural: a secondary inspection model requires a functional graph. A subclassed model whose call is imperative Python has no symbolic intermediate tensors to expose, so its call method must record the requested values explicitly.

Nothing here needs detaching or removing: the captured intermediates are ordinary arrays with no autograd graph attached, the dictionary exists for this one call, and no observer stays registered on the model. The method_outputs=nnx.Intermediate argument records every method return. For finer control, a module can call self.sow(nnx.Intermediate, 'name', value) at selected points, and nnx.capture(model, nnx.Intermediate) then returns only those named values.

Two rules keep hooks safe. First, detach before you stash: the hook above reduces the output to a Python float on the spot; a hook that stores arrays should store output.detach() (add .copy() if the array may later be updated in place), because under autograd.record a stashed output keeps the computation graph of every forward pass alive until the list is cleared. Second, keep the handle and call handle.detach(): register_forward_hook returns a HookHandle, and the hook stays registered for the block’s lifetime otherwise (the handle also works as a context manager that detaches on exit). One caveat specific to Gluon: hooks live in the Python __call__ wrapper, so hybridize() warns that hooks on the children of a hybridized block will not fire; keep a block unhybridized while observing it.

5.8.3.2 A NaN Finder

When a loss becomes NaN at step 40,000, the useful question is which layer produced the first non-finite value. Checking every leaf module’s output for finiteness can identify the first affected layer in one forward pass. The example injects a NaN into a weight deep in the stack:

def make_finite_check(name):
    def check(module, inputs, output):
        if not torch.isfinite(output).all():
            raise RuntimeError(f'first non-finite output in {name}')
    return check

handles = [m.register_forward_hook(make_finite_check(name))
           for name, m in net.named_modules()
           if len(list(m.children())) == 0]
with torch.no_grad():
    net[3].body[0].weight[0, 0] = float('nan')  # sabotage one layer
try:
    net(torch.randn(2, 20))
except RuntimeError as e:
    print(e)
for h in handles:
    h.remove()

first non-finite output in 3.body.0

class ResidualBlock(tf.keras.layers.Layer):
    def __init__(self, num_hiddens, **kwargs):
        super().__init__(**kwargs)
        self.body = tf.keras.Sequential([
            tf.keras.layers.Dense(num_hiddens, activation='relu'),
            tf.keras.layers.Dense(num_hiddens)])

    def call(self, X):
        return X + self.body(X)

class Checked(tf.keras.Model):
    def __init__(self, layers):
        super().__init__()
        self.seq, self.first_bad = layers, None  # a plain list is tracked

    def call(self, X):
        for layer in self.seq:
            X = layer(X)
            if self.first_bad is None and not bool(
                    tf.reduce_all(tf.math.is_finite(X))):
                self.first_bad = layer.name
        return X

tf.keras.utils.set_random_seed(0)
checked = Checked([tf.keras.layers.Dense(64)]
                  + [ResidualBlock(64, name=f'block{k}') for k in range(1, 9)]
                  + [tf.keras.layers.Dense(10)])
_ = checked(tf.random.normal((2, 20)))  # build the weights
kernel = checked.seq[3].body.layers[0].kernel
kernel[0, 0].assign(float('nan'))  # sabotage one layer
_ = checked(tf.random.normal((2, 20)))
print('first non-finite output in', checked.first_bad)

first non-finite output in block3

saved = net.layers[3].body.layers[0].kernel[0, 0]
net.layers[3].body.layers[0].kernel[0, 0] = float('nan')
_, inter = nnx.capture(
    net, nnx.Intermediate, method_outputs=nnx.Intermediate)(X)

def get_path(tree, path):
    for key in path:
        tree = tree[key]
    return tree

for path, module in nnx.iter_modules(net):
    if isinstance(module, nnx.Linear):
        out = get_path(inter, path)['__call__'][0]
        if not jnp.isfinite(out).all():
            print('first non-finite output in', path)
            break

net.layers[3].body.layers[0].kernel[0, 0] = saved  # net is shared; undo the sabotage

first non-finite output in ('layers', 3, 'body', 'layers', 0)

def make_finite_check(name):
    def check(block, inputs, output):
        if not np.isfinite(output).all():
            raise RuntimeError(f'first non-finite output in {name}')
    return check

leaves = ([('0', net[0])]
          + [(f'{k}.body.{j}', net[k].body[j])
             for k in range(1, 9) for j in (0, 1)]
          + [('9', net[9])])
handles = [blk.register_forward_hook(make_finite_check(name))
           for name, blk in leaves]
net[3].body[1].weight.data()[0, 0] = float('nan')  # sabotage one layer
try:
    net(np.random.normal(0, 1, (2, 20)))
except RuntimeError as e:
    print(e)
for h in handles:
    h.detach()

first non-finite output in 3.body.1

The report names module 3.body.0, the layer containing the injected NaN, without requiring manual per-layer print statements.

The report names block3, the layer containing the injected NaN, without requiring manual per-layer print statements. The check runs on every forward pass until you edit it out of call, the price of building observation into the model rather than attaching it from outside. For temporary diagnosis, TensorFlow also provides this behavior through a switch: tf.debugging.enable_check_numerics() instruments every operation and reports the first one to produce an inf or NaN.

We inspect the captured outputs of linear modules in object-graph order, which matches execution order for this sequential network. The report names ('layers', 3, 'body', 'layers', 0), the layer containing the injected NaN, without requiring manual per-layer print statements.

The report names block 3.body.1, the layer containing the injected NaN, without requiring manual per-layer print statements. Two Gluon-specific details appear in the setup. First, blocks carry no path names and there is no named_modules equivalent, so the example constructs the list of leaves explicitly. Second, it injects the NaN into the block’s second dense layer: Gluon’s fused Dense(..., activation='relu') kernel maps NaN to \(0\), so a NaN before the fused activation can be replaced before a hook observes it. Diagnostic hooks must therefore be placed where the target condition remains observable. Registration order does not matter: hooks fire in execution order, so the first completed forward with a non-finite output identifies the affected block.

5.8.3.3 Backward Hooks and Beyond

Gradients get the same treatment. module.register_full_backward_hook(f) runs f(module, grad_input, grad_output) after the module’s gradients are computed, which is how you log activation-gradient norms or catch an explosion at its source. A returned replacement changes grad_input; it does not edit the module’s parameter gradients. Use this API rather than the older register_backward_hook, which is deprecated and unreliable for modules with multiple inputs. Like forward hooks, it returns a handle to remove and should detach anything it stores. For the specific job of extracting intermediate features from a standard vision backbone, torchvision’s create_feature_extractor is the production upgrade: it traces the model with torch.fx and returns a module that outputs the requested internal nodes directly, provided the model is traceable, whereas hooks work on anything.

There are no backward hooks, but gradients do not need them: tf.GradientTape already hands back the gradient of every variable as a value. To log per-layer gradient norms, catch exploding gradients at their source, or experiment with per-layer clipping, compute the gradients and inspect or transform them like any other data. For extracting features from a pretrained backbone, a secondary functional model provides the required output: tf.keras.Model(backbone.input, backbone.get_layer('avg_pool').output) turns any functional backbone into a feature extractor by naming the tensor you want.

Gradients need no hook at all, because they are not events that fire inside a backward pass: nnx.grad returns them as a tree of values in the same shape as the parameters. To log per-layer gradient norms, catch exploding gradients at their source, or experiment with per-layer clipping, compute the gradient tree and inspect or transform it like any other data:

Gluon has no backward hooks: nothing can be attached to observe gradients as the backward pass computes them. What it has is post-hoc inspection, and for most debugging that is enough: after loss.backward() every parameter’s gradient sits in param.grad(), so to log per-layer gradient norms or catch an explosion you iterate net.collect_params() and read them like any other array. What post-hoc inspection cannot do is act during the pass, so per-layer gradient clipping in the style of exercise 3 is written into the training step instead; Gluon’s Trainer splits step into allreduce_grads and update precisely so that such code can stand between them.

with tf.GradientTape() as tape:
    loss = tf.reduce_mean(net(X) ** 2)
grads = tape.gradient(loss, net.trainable_variables)
for v, g in list(zip(net.trainable_variables, grads))[10:12]:
    print(v.path, float(tf.norm(g)))  # block 3's first layer

sequential_5/dense_11/kernel 13.867950439453125
sequential_5/dense_11/bias 2.8245253562927246

grads = nnx.grad(lambda model: (model(X) ** 2).mean())(net)
norms = jax.tree_util.tree_map(jnp.linalg.norm, grads)
print(norms['layers'][3]['body']['layers'][0])

State({
  'bias': Param( # 1 (4 B)
    value=Array(1.9621452, dtype=float32)
  ),
  'kernel': Param( # 1 (4 B)
    value=Array(15.665125, dtype=float32)
  )
})

5.8.4 Summary

Reproducibility requires an inventory of random streams and arithmetic: initialization, dropout, shuffling, augmentation, and loader workers each use a generator, while deterministic kernels are required for bitwise repetition.

torch.manual_seed covers torch’s CPU and CUDA generators; NumPy and random need their own seeds in the main process. DataLoader derives distinct worker seeds and a supplied generator pins the shuffle order; worker_init_fn also configures generator objects that the loader cannot discover. Seeding fixes random streams but does not determine arithmetic. torch.use_deterministic_algorithms(True) additionally requests deterministic implementations and raises on operations that cannot comply.

tf.keras.utils.set_random_seed covers Python’s random, NumPy, and TensorFlow’s global generator in one call; tf.random.Generator objects and Dataset.shuffle(seed=...) carry their seeds explicitly. Seeding fixes random streams but does not determine arithmetic. tf.config.experimental.enable_op_determinism() additionally requests deterministic implementations, raising on operations that cannot comply; enable it near program startup, before creating operations.

The JAX PRNG portion of the inventory is explicit: the same key gives the same draws, and independent JAX streams come from jax.random.split. A whole program must still inventory Python random, NumPy, data loaders, host callbacks, and external libraries. Keys make JAX draws repeatable but do not determine arithmetic. The CPU operations tested in this chapter were deterministic in this environment; on GPU, --xla_gpu_deterministic_ops=true requests deterministic kernels.

np.random.seed covers MXNet’s generator on every device in one call; NumPy’s legacy global generator and Python’s random need their own seeds, and Gluon’s loader consults NumPy for its shuffle order while offering no per-worker seeding hook, so worker randomness must be pinned inside the dataset. Seeding fixes random streams but does not determine arithmetic. MXNet has no determinism switch, only the autotuning control MXNET_CUDNN_AUTOTUNE_DEFAULT=0.

Bitwise agreement is a debugging tool; conclusions that hold across seeds are the scientific goal.

For looking inside a model, forward and backward hooks observe any module without editing it, subject to two rules: detach what you stash, and remove the handle.

For looking inside a model there is no hook to attach: declare the tensors you want as extra outputs of a functional model, or override call where you own the code, and read gradients as values from tf.GradientTape.

For looking inside a model, nnx.capture(..., method_outputs=...) returns submodule outputs from an unmodified call, sow records named values from the inside, and gradients are values from nnx.grad you inspect directly.

For looking inside a model, forward hooks and pre-hooks observe any block without editing it, subject to two rules: detach what you stash, and detach the HookHandle when done. Backward hooks do not exist; gradients are read after the fact from param.grad().

5.8.5 Exercises

Extend train_once to load data through your framework’s parallel data loader (e.g., PyTorch’s DataLoader, with num_workers=4). Give the dataset its own np.random.default_rng() object and use it to add noise as each example loads. On a process-based loader, check whether workers copied the same generator state. Using the per-worker seeding hook this section introduced for your framework, replace the dataset’s stashed generator with one derived from the worker’s own seed, then verify agreement across two runs.
Write a forward hook (or your framework’s equivalent introduced in this section) that counts multiply-accumulate operations for every linear layer from the shapes of its input and weight, and use it to report per-layer and total FLOPs for the residual stack above. Check the total against a hand count.
Using a backward hook (or your framework’s equivalent introduced in this section), record the norm of each block’s output gradient. After the backward pass, compare these activation-gradient norms with the parameter-gradient norms in the same block. Which one first reveals an exploding backward signal?
A forward hook that returns a value replaces the module’s output. Use one to zero out the output of a single residual block’s body (turning the block into the identity) and measure how the network’s output changes. Which block matters most, and how would you find out with one loop?

Rebuild the activation-statistics table two more ways: with a Checked-style call override that stashes tf.math.reduce_std of every layer’s output, and on a model you did not write, a tf.keras.applications backbone, by naming layers with get_layer. Compare the three contracts introduced here: functional inspection models, call overrides, and PyTorch-style mutable hooks: which requires a symbolic graph, which requires owning the model’s code, and which black-box models does each admit?

Rebuild the activation-statistics table by editing ResidualBlock to call self.sow(nnx.Intermediate, 'body_out', ...) on its body’s output, then collect those values with nnx.capture(net, nnx.Intermediate). Compare capture-all-methods, opt-in sow, and PyTorch-style mutable hooks: which requires touching model code, which can retain tensors or hook state, and which is applicable to a model you do not own?

Rebuild the activation-statistics table two more ways: with register_forward_pre_hook, recording the standard deviation of each block’s input (the pre-hook receives the argument tuple, so unpack it), and by reading the source of Block.summary(), which builds its whole table from one forward hook, then calling it on the residual stack. Then revisit exercise 4: Gluon ignores a hook’s return value, so an output-replacing hook cannot be written; perform the same ablation by editing the block’s forward instead, and compare the two contracts, observe-only hooks against PyTorch’s mutable ones, for a model you own and for one you do not.