week01_notebook.ipynb. Source:
view on GitHub.
lax.scan, fori_loop, jit, vmap, and associative scan to dynamical-systems habits.Week 1 is the syntax foundation. The same recurrence you would write in a loop becomes a pure step function plus an explicit carry. Once that pattern is clear, SSM layers become less mysterious: they are mostly structured recurrences, vectorized batches, and scans over sequence length.
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/week01/primitives.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 prepares the notebook environment. The path logic finds shared plotting helpers from either the repo root or the notebook directory, then imports Week 1 primitives. Notice that the numerical functions live in experiments/jax/week01/primitives.py; the notebook is here to explain and exercise them rather than hide them.
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 sys
from pathlib import Path
# 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=1)
import jax, jax.numpy as jnp
import matplotlib.pyplot as plt
from primitives import (
damped_oscillator_scan,
damped_oscillator_loop,
damped_oscillator_fori,
oscillator_endpoint_jit,
benchmark_oscillator_implementations,
coupled_oscillators_scan,
coupled_oscillators_jacobian,
vmap_oscillator_ensemble,
vmap_coupled_oscillator_jacobians,
lyapunov_exponents_qr,
cumsum_associative,
ema_sequential,
ema_associative,
)
print(f"JAX {jax.__version__} | {jax.default_backend()} | Experiments: {EXPERIMENT_DIR}")
The first computation runs the damped oscillator with lax.scan. The carry is the oscillator state and each scan step advances one time step. Check the printed shape first: sequence length is the leading axis, and the two state coordinates are position and velocity.
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.
omega, gamma, dt, n_steps = 2.0, 0.5, 0.01, 1000
states_scan = damped_oscillator_scan(1.0, 0.0, omega, gamma, dt, n_steps)
states_loop = damped_oscillator_loop(1.0, 0.0, omega, gamma, dt, n_steps)
x, v = states_scan[:, 0], states_scan[:, 1]
t = jnp.arange(n_steps) * dt
energy = 0.5 * (v**2 + omega**2 * x**2)
print(f"scan vs loop max diff: {jnp.max(jnp.abs(states_scan - states_loop)):.2e}")
grad_gamma = jax.grad(lambda g: damped_oscillator_scan(1.0, 0.0, 2.0, g, 0.01, 500)[-1, 0])(0.5)
print(f"d(x_final)/d(gamma) = {grad_gamma:.4f}")
This plotting cell converts the scan output into the guide figure. The important reading habit is to connect the visual decay with the recurrence: each point is not independently predicted; it is produced by repeatedly applying the same step function.
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.
fig, (ax1, ax2) = create_tufte_figure(1, 2, figsize=(8, 3.5))
ax1.plot(t, x, label='$x(t)$', lw=1.2, color=C_ACCENT)
ax1.plot(t, v, label='$v(t)$', lw=1.2, alpha=0.7, color=C_ALERT)
set_tufte_title(ax1, 'Damped oscillator')
set_tufte_labels(ax1, 'Time', 'State')
ax1.legend(frameon=False, fontsize=8)
ax2.plot(t, energy, color=C_ALERT, lw=1.2)
set_tufte_title(ax2, 'Energy decay')
set_tufte_labels(ax2, 'Time', 'Energy')
ax2.set_yscale('log')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex1_energy_decay')
plt.show()
fori_loop is the scalar loop counterpart to scan. It returns only the final carry, so it is useful when you do not need the whole trajectory. The guided modify-and-run prompt is to change n_steps and confirm that both loop styles land at the same final state.
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.
scan_final = damped_oscillator_scan(1.0, 0.0, 2.0, 0.5, 0.01, 500)[-1]
fori_final = damped_oscillator_fori(1.0, 0.0, 2.0, 0.5, 0.01, 500)
jit_final = oscillator_endpoint_jit(1.0, 0.0, 2.0, 0.5, 0.01, 500)
print(f"scan: {scan_final} fori: {fori_final} jit: {jit_final}")
print(f"fori diff: {jnp.max(jnp.abs(scan_final - fori_final)):.2e} jit diff: {jnp.max(jnp.abs(scan_final - jit_final)):.2e}")
This benchmark introduces jit. The first call includes compilation overhead; repeated calls measure the compiled computation. Expected observation: compiled JAX code wins when the loop is long enough for compilation cost to be amortized.
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.
print("Timing benchmark (3 repeats, 10000 steps)...")
timings = benchmark_oscillator_implementations(n_steps=10000, n_repeats=3)
for name, t_val in timings.items():
print(f" {name:12s}: {t_val*1000:.2f} ms")
# Log-scale dot plot with direct labels — shows order-of-magnitude differences
# that bar charts compress. Each row is a method; position on x-axis shows timing.
names = list(timings.keys())
values = [v * 1000 for v in timings.values()] # Convert to ms
colors = [C_ALERT, C_ACCENT, C_HIGHLIGHT]
fig, ax = create_tufte_figure(figsize=(6, 2.5))
y_pos = range(len(names))
ax.scatter(values, y_pos, s=80, c=colors, zorder=5, edgecolors='white', lw=0.5)
ax.set_yticks(list(y_pos))
ax.set_yticklabels(names)
ax.set_xscale('log')
set_tufte_labels(ax, xlabel='Wall-clock time (ms, log scale)')
set_tufte_title(ax, '10k-step oscillator: compilation eliminates interpreter overhead')
# Direct labels with actual values
for i, (name, val) in enumerate(zip(names, values)):
direct_label(ax, val, i, f'{val:.1f} ms', offset=(8, -3), fontsize=8)
ax.invert_yaxis() # Fastest at top
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex1b_timing_comparison')
plt.show()
The coupled oscillator example scales the state from one oscillator to several interacting oscillators. The key syntax is still a scan carry, but the state vector is longer and the update function includes a coupling matrix.
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_osc = 4
state0 = jnp.zeros(2 * n_osc).at[0].set(1.0)
states = coupled_oscillators_scan(state0, omega=1.0, coupling=0.1, dt=0.01, n_steps=1000)
J = coupled_oscillators_jacobian(state0, omega=1.0, coupling=0.1, dt=0.01)
eigenvalues = jnp.linalg.eigvals(J)
print(f"Jacobian: {J.shape}, eigenvalues inside unit disk: {jnp.all(jnp.abs(eigenvalues) <= 1.0 + 1e-6)}")
The plot and Jacobian view teach local sensitivity. You are not just plotting trajectories; you are asking how a small perturbation would propagate through the update map.
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.
fig, (ax1, ax2) = create_tufte_figure(1, 2, figsize=(8, 3.5))
for i in range(n_osc):
ax1.plot(jnp.arange(1000)*0.01, states[:, i], lw=1.0)
set_tufte_title(ax1, f'Coupled oscillators ($N={n_osc}$)')
set_tufte_labels(ax1, 'Time', 'Position')
theta = jnp.linspace(0, 2*jnp.pi, 100)
ax2.plot(jnp.cos(theta), jnp.sin(theta), color=C_BASELINE, ls='--', alpha=0.5)
ax2.scatter(jnp.real(eigenvalues), jnp.imag(eigenvalues), c=C_ACCENT, s=60, zorder=5, edgecolors='white', lw=0.5)
set_tufte_title(ax2, 'Jacobian eigenvalues')
set_tufte_labels(ax2, r'Re($\lambda$)', r'Im($\lambda$)')
ax2.set_aspect('equal')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex2_jacobian_eigenvalues')
plt.show()
vmap turns a single-trajectory function into a batched function without manually writing a Python loop. The expected observation is that the output grows a new leading batch axis while the per-example logic stays unchanged.
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.
ensemble = vmap_oscillator_ensemble(jnp.array([1.0, 0.5, 0.1, -0.5]), jnp.zeros(4), 2.0, 0.5, 0.01, 500)
print(f"Ensemble: {ensemble.shape}")
pair = vmap_oscillator_ensemble(jnp.array([1.0, 2.0]), jnp.zeros(2), 2.0, 0.0, 0.001, 500)
print(f"Linearity (undamped): |2*traj[0] - traj[1]| = {jnp.max(jnp.abs(2*pair[0] - pair[1])):.2e}")
traj_v = coupled_oscillators_scan(jnp.zeros(8).at[0].set(1.0), 1.0, 0.1, 0.01, 100)
jacs = vmap_coupled_oscillator_jacobians(traj_v[::25], 1.0, 0.1, 0.01)
print(f"Batched Jacobians: {jacs.shape}, variation: {jnp.max(jnp.abs(jacs - jacs[0])):.2e}")
This Lyapunov cell uses QR reorthogonalization. The syntax looks more advanced, but the concept is familiar from dynamical systems: evolve tangent vectors, periodically re-normalize them, and accumulate expansion rates.
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.
A_diag = jnp.diag(jnp.array([-1.0, -2.0]))
dynamics_diag = lambda s: A_diag @ s
exp_diag = lyapunov_exponents_qr(dynamics_diag, dim=2, dt=0.01, n_steps=10000)
print(f"Computed: {exp_diag} Expected: [-1, -2] Sum: {jnp.sum(exp_diag):.2f} = trace: {jnp.trace(A_diag):.0f}")
A_off = jnp.array([[-1.0, 0.3], [-0.3, -2.0]])
exp_off = lyapunov_exponents_qr(lambda s: A_off @ s, dim=2, dt=0.01, n_steps=10000)
print(f"Off-diag: {exp_off} Sum: {jnp.sum(exp_off):.2f} = trace: {jnp.trace(A_off):.0f}")
The convergence experiment shows why QR-based Lyapunov estimates need time. Modify the step counts and watch the estimate settle; short trajectories can be noisy even when the system is simple.
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.
steps_range = [500, 1000, 2000, 5000, 10000, 20000]
exp_conv = jnp.array([lyapunov_exponents_qr(dynamics_diag, 2, 0.01, ns) for ns in steps_range])
fig, ax = create_tufte_figure(figsize=(6, 3.5))
ax.plot(steps_range, exp_conv[:, 0], 'o-', lw=1.2, color=C_ACCENT)
ax.plot(steps_range, exp_conv[:, 1], 's-', lw=1.2, color=C_ALERT)
ax.axhline(-1.0, color=C_ACCENT, ls='--', alpha=0.3)
ax.axhline(-2.0, color=C_ALERT, ls='--', alpha=0.3)
set_tufte_title(ax, 'Lyapunov exponent convergence via QR')
set_tufte_labels(ax, 'QR steps', 'Lyapunov exponent')
ax.set_xscale('log')
direct_label_line(ax, steps_range, exp_conv[:, 0], '$\\lambda_1 \\to -1$', position='end', fontsize=8)
direct_label_line(ax, steps_range, exp_conv[:, 1], '$\\lambda_2 \\to -2$', position='end', fontsize=8)
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex3_lyapunov_convergence')
plt.show()
The cumulative-sum warm-up is the smallest associative-scan example. Addition is associative, so prefix sums can be parallelized. This is the same computational idea that later makes S5 and Mamba scans possible.
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.
test_x = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(f"assoc_scan: {cumsum_associative(test_x)} jnp.cumsum: {jnp.cumsum(test_x)} match: {jnp.allclose(cumsum_associative(test_x), jnp.cumsum(test_x))}")
x_rand = jax.random.normal(jax.random.PRNGKey(99), (256,))
print(f"Random L=256 max diff: {jnp.max(jnp.abs(cumsum_associative(x_rand) - jnp.cumsum(x_rand))):.2e}")
EMA is an affine recurrence, not just an addition. The notebook compares sequential and associative implementations so you can see that affine composition is the real generalization.
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.
x_ema = jax.random.normal(jax.random.PRNGKey(42), (100,))
alpha = 0.1
ema_seq = ema_sequential(x_ema, alpha)
ema_assoc = ema_associative(x_ema, alpha)
print(f"Sequential vs associative max diff: {jnp.max(jnp.abs(ema_seq - ema_assoc)):.2e}")
The final plot makes the scan abstraction concrete. The input is noisy, the EMA is smoothed, and both sequential and associative versions agree. Expected observation: the curves overlap; if they do not, the associative operator is wrong.
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.
fig, (ax1, ax2) = create_tufte_figure(1, 2, figsize=(8, 3.5))
ax1.plot(x_ema, alpha=0.3, color=C_BASELINE, lw=0.8, label='Input')
ax1.plot(ema_seq, lw=1.5, color=C_ACCENT)
ax1.plot(ema_assoc, lw=1.5, ls='--', color=C_ALERT)
set_tufte_title(ax1, f'EMA ($\\alpha={alpha}$)')
set_tufte_labels(ax1, 'Step', '$y_t$')
direct_label(ax1, 99, float(ema_seq[-1]), 'Sequential', offset=(5, 3), fontsize=8, color=C_ACCENT)
direct_label(ax1, 99, float(ema_assoc[-1]), 'Associative', offset=(5, -10), fontsize=8, color=C_ALERT)
ax2.semilogy(jnp.abs(ema_seq - ema_assoc), color=C_HIGHLIGHT, lw=1.0)
set_tufte_title(ax2, 'Numerical agreement')
set_tufte_labels(ax2, 'Step', '|seq - assoc|')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex4_ema_comparison')
plt.show()
This section makes the notebook intentionally self-contained. The week topic is JAX Primitives Through Dynamical Systems, and the working skill is turning Python loops into compiled array programs. 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/week01/primitives.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/week01/primitives.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/week01/ -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 2 moves from primitives to trainable modules. The state-carry idea stays, but parameters now live inside Flax NNX objects and optimization updates them.