week02_notebook.ipynb. Source:
view on GitHub.
Week 2 is the engineering bridge between mathematical primitives and trainable experiments. The important move is not that the model is deep; it is that every later week can reuse the same shape conventions, logging pattern, optimizer boundary, and synthetic task interface.
This notebook is intentionally standalone. The PDF chapter gives a more polished narrative, but you should be able to run this notebook by itself and still understand what the computation is doing. The repeated pattern is: name the mathematical object, inspect the shape contract, run the smallest useful example, then connect the result back to the production implementation at experiments/jax/common/harness.py.
Read each markdown block before running the next code cell. When a block says Guided modify-and-run, make the suggested small edit, rerun just that cell, and compare your observation with the inline Expected observation. These prompts are deliberately small because the goal is code literacy: you should know which axis is batch, which axis is time, where parameters live, and what failure would look like before you move to larger experiments.
This setup cell imports Flax NNX, Optax, the shared harness, and plotting helpers. The key idea is that Week 2 moves from pure functions to stateful modules, but NNX keeps the state explicit and inspectable.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
import os, sys
from pathlib import Path
# GPU memory management for 8GB cards (RTX 2070S)
os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.70")
# Add shared utilities to path
for p in ["../shared", "guides/shared"]:
resolved_path = Path(p).resolve()
if (resolved_path / "plot_utils.py").exists():
if str(resolved_path) not in sys.path:
sys.path.insert(0, str(resolved_path))
break
from plot_utils import (
setup_notebook, save_figure, create_tufte_figure, apply_tufte_style,
add_subtle_grid, set_tufte_title, set_tufte_labels, direct_label,
direct_label_line, SSM_COLORS, C_ACCENT, C_ALERT, C_HIGHLIGHT, C_BASELINE,
)
EXPERIMENT_DIR, FIGURE_DIR = setup_notebook(week=2)
import jax, jax.numpy as jnp, matplotlib.pyplot as plt
import flax
from flax import nnx
import optax
from experiments.jax.common.harness import (
SimpleSeqModel, make_noisy_sine_batch, train_step,
ResidualBlock, StackedModel,
SSMParam, ExperimentHarness, PrintLogger,
)
print(f"JAX {jax.__version__} | Flax {flax.__version__} | {jax.default_backend()}")
print(f"Experiments: {EXPERIMENT_DIR}")
The first model call checks shape discipline. SimpleSeqModel(features=64) maps (batch, length, channels) to the same output shape. Modify the feature count and confirm that shape stays stable while parameter count changes.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
model = SimpleSeqModel(features=64, rngs=nnx.Rngs(0))
x_demo = jnp.ones((2, 128, 1))
y_demo = model(x_demo)
print(f"Input: {x_demo.shape} -> Output: {y_demo.shape}")
print(f"Parameters: {sum(p.size for p in jax.tree.leaves(nnx.state(model, nnx.Param)))}")
The training-step cell introduces optimizer state. Read the function signature carefully: model, optimizer, inputs, and targets are explicit values. Expected observation: loss is finite and gradients update trainable parameters.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
model = SimpleSeqModel(features=64, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
losses = []
key = jax.random.key(0)
for step in range(300):
key, subkey = jax.random.split(key)
x, y = make_noisy_sine_batch(subkey, batch_size=32, seq_len=128)
loss = train_step(model, optimizer, x, y)
losses.append(float(loss))
print(f"Loss: {losses[0]:.4f} -> {losses[-1]:.4f} (ratio: {losses[-1]/losses[0]:.2f})")
The noisy-sine example is the reusable synthetic task. The point is not the task difficulty; it is that later SSM layers can be compared under the same data-generation and evaluation contract.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# Show data + predictions
key = jax.random.key(42)
x_vis, y_vis = make_noisy_sine_batch(key, batch_size=4, seq_len=128, noise_std=0.1)
preds = model(x_vis) # model already trained above
fig, axes = create_tufte_figure(1, 2, figsize=(10, 4))
# Loss curve
axes[0].plot(losses, color=C_ACCENT, lw=0.8, alpha=0.5)
window = 20
if len(losses) >= window:
ma = [sum(losses[i:i+window])/window for i in range(len(losses)-window+1)]
axes[0].plot(range(window-1, len(losses)), ma, color=C_ACCENT, lw=2)
set_tufte_title(axes[0], 'Training loss')
set_tufte_labels(axes[0], 'Step', 'MSE loss')
# Predictions vs targets (first sample)
t = jnp.arange(128)
axes[1].plot(t, x_vis[0, :, 0], color=C_BASELINE, alpha=0.4, lw=0.8)
axes[1].plot(t, y_vis[0, :, 0], color=C_ACCENT, lw=1.5)
axes[1].plot(t, preds[0, :, 0], color=C_ALERT, lw=1.5, ls='--')
set_tufte_title(axes[1], 'Noisy sine prediction')
set_tufte_labels(axes[1], 'Time step', 'Value')
direct_label(axes[1], 127, float(y_vis[0, -1, 0]), 'Target', offset=(5, 3), fontsize=8, color=C_ACCENT)
direct_label(axes[1], 127, float(preds[0, -1, 0]), 'Prediction', offset=(5, -10), fontsize=8, color=C_ALERT)
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex3_loss_and_predictions')
plt.show()
Layer stacking adds depth while keeping the sequence interface fixed. The guided modify-and-run prompt is to compare shallow and deeper models and ask whether extra capacity improves convergence or just adds parameters.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# Compare training trajectories at different depths
# Overlaid loss curves show *how* convergence differs, not just the endpoint
layer_counts = [1, 2, 4, 8]
all_curves = {}
colors_depth = [C_ACCENT, C_ALERT, C_HIGHLIGHT, C_BASELINE]
for n_layers in layer_counts:
sm = StackedModel(features=32, num_layers=n_layers, rngs=nnx.Rngs(0))
harness = ExperimentHarness(sm, total_steps=200, backbone_lr=1e-3,
logger=PrintLogger(every=999))
curve = harness.train(make_noisy_sine_batch, n_steps=200, batch_size=16)
all_curves[n_layers] = curve
final = sum(curve[-10:]) / 10
print(f" {n_layers} layers: final_loss={final:.4f}")
fig, ax = create_tufte_figure(figsize=(7, 4))
window = 10
for (n_layers, curve), c in zip(all_curves.items(), colors_depth):
# Smoothed loss curve
if len(curve) >= window:
smoothed = [sum(curve[i:i+window])/window for i in range(len(curve)-window+1)]
steps = list(range(window-1, len(curve)))
ax.plot(steps, smoothed, lw=1.5, color=c)
# Direct label at endpoint
final_val = sum(curve[-10:]) / 10
direct_label_line(ax, steps, smoothed,
f'{n_layers}L ({final_val:.3f})',
position='end', fontsize=8, color=c)
set_tufte_title(ax, 'Deeper models converge faster and lower')
set_tufte_labels(ax, 'Training step', 'MSE loss (smoothed)')
ax.set_yscale('log')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex4_layer_scaling')
plt.show()
The harness cell packages repeated experiment ceremony. This is acceptable abstraction because the notebook first showed the underlying pieces: model construction, batch generation, optimizer, and metric logging.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
n_steps = 500
# Single optimizer: fixed LR, no schedule
model_single = SimpleSeqModel(features=64, rngs=nnx.Rngs(0))
opt_single = nnx.Optimizer(model_single, optax.adam(1e-3), wrt=nnx.Param)
losses_single = []
key = jax.random.key(0)
for step in range(n_steps):
key, subkey = jax.random.split(key)
x, y = make_noisy_sine_batch(subkey, batch_size=32)
loss = train_step(model_single, opt_single, x, y)
losses_single.append(float(loss))
# Dual optimizer harness: warmup + cosine + grad clip
model_dual = SimpleSeqModel(features=64, rngs=nnx.Rngs(0))
harness = ExperimentHarness(model_dual, total_steps=n_steps, backbone_lr=1e-3,
logger=PrintLogger(every=999))
losses_dual = harness.train(make_noisy_sine_batch, n_steps=n_steps, batch_size=32)
window = 20
fig, ax = create_tufte_figure(figsize=(7, 4))
def smooth(vals, w):
return [sum(vals[i:i+w])/w for i in range(len(vals)-w+1)]
steps_sm = list(range(window-1, len(losses_single)))
ax.plot(steps_sm, smooth(losses_single, window), color=C_ALERT, lw=1.5)
ax.plot(steps_sm, smooth(losses_dual, window), color=C_ACCENT, lw=1.5)
set_tufte_title(ax, 'Warmup + cosine decay vs fixed learning rate')
set_tufte_labels(ax, 'Step', 'MSE loss (smoothed)')
# Direct labels
direct_label_line(ax, steps_sm, smooth(losses_single, window),
f'Fixed LR ({sum(losses_single[-10:])/10:.4f})',
position='end', fontsize=8, color=C_ALERT)
direct_label_line(ax, steps_sm, smooth(losses_dual, window),
f'Harness ({sum(losses_dual[-10:])/10:.4f})',
position='end', fontsize=8, color=C_ACCENT)
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex5_dual_vs_single')
plt.show()
print(f"Single: final={sum(losses_single[-10:])/10:.4f}")
print(f"Harness: final={sum(losses_dual[-10:])/10:.4f}")
This section makes the notebook intentionally self-contained. The week topic is Flax NNX Modules and Training Loop, and the working skill is keeping model state, optimizer state, and random state explicit. The PDF may give a smoother chapter narrative, but the notebook should still answer four practical questions on its own: what object are you building, what syntax expresses it, what shape proves the code is wired correctly, and what numerical observation would make you suspicious.
For this week, keep three layers separate. First is the mathematical object: a recurrence, table, matrix, discretization, filter, or evidence record. Second is the implementation syntax: the JAX transform, Flax NNX state object, FFT call, dataclass, or array operation that expresses the object. Third is the validation habit: shape checks, equality against a reference implementation, stability checks, or a compact evidence ledger. When notebooks feel abstract, it is usually because these layers are collapsed into one imported function call. The markdown here deliberately pulls them apart before the code runs.
A good reading pass is slow and mechanical. Before a code cell, predict the output shape and the one numerical property that should hold. After the cell, check that prediction before looking at plots or final losses. If a plot is produced, ask what computation created every axis. If a scalar comparison is printed, ask which reference implementation it is comparing against. If a table row is produced, ask what would falsify it. This is the same discipline you will need when debugging the production path in experiments/jax/common/harness.py.
Use one-edit experiments. Change one sequence length, state dimension, step size, feature count, decay rate, or row in a reading table. Do not change three things at once. The expected observation should be simple enough to say before rerunning: the leading axis should become longer, the output channel count should stay fixed, the stability magnitude should move toward zero, or the equivalence error should remain near machine precision. If the result violates that expectation, the notebook has done its job: it found the exact boundary where your mental model was incomplete.
For every major cell in this notebook, you should be able to fill in this checklist: input shape, output shape, state carried across time, parameter or evidence object being changed, reference result if one exists, and the failure mode that would matter. The answers are included inline because this is a self-teaching notebook. Use them actively: read the answer, rerun the cell, then explain why the answer follows from the code rather than memorizing the printed value.
Pause before the summary and explain the notebook aloud in three passes. First, describe the mathematical or methodological object without code. Second, describe the exact array or Python object that represents it. Third, describe the validation signal printed by the notebook. This deliberate repetition is useful because later weeks reuse the same ideas with more dimensions and less scaffolding.
A strong answer should include at least one shape statement and one failure statement. For example: "the sequence axis is length L, the state axis is size N, and the equivalence check would fail if the scan composition were ordered incorrectly." For reading weeks, replace shape with schema: "the evidence row has a family, claim, limitation, and follow-up, and the failure would be mixing evidence with opinion." This habit turns notebook reading into implementation readiness.
When you adapt the code, prefer small edits that preserve the contract. If a small edit breaks the notebook, do not immediately patch around the error. Read the message, identify the violated shape or state assumption, and then decide whether the production implementation should prevent that violation with validation or whether the notebook simply taught you an expected boundary.
You should now be able to state the week in code terms: what object is being represented, what shape each major tensor carries, which production file contains the durable implementation, and what observation confirms the implementation is behaving correctly. The notebook uses experiments/jax/common/harness.py as the reference path; if you want to go deeper, read that file next to the tests rather than treating the import as magic.
Run pytest experiments/jax/week02/ -v after changing the implementation or adapting an example. For notebook-only edits, also run make notebook-check so sparse markdown, committed outputs, and missing syntax supplements fail immediately.
Week 3 steps back from implementation and asks which operator families matter, why they exist, and what evidence supports or weakens each one.