week11_notebook.ipynb. Source:
view on GitHub.
Week 11 shows a different route to long context. Instead of a structured recurrence, Hyena uses long causal filters and FFT convolution. The implementation is short, but the shape discipline is strict: batch, sequence, channel, filter, and feedthrough must agree.
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/week11/hyena_lineage.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.
The production module exposes fftconv and a naive reference. The reference is slow but obvious; the FFT version is fast but easy to get wrong if padding or axes are wrong. This notebook starts with a tiny input so the comparison is interpretable.
Guided modify-and-run. Change L from 8 to 16 or change the kernel to a lag-one filter. The output shape should remain (B, L, D).
Expected observation. FFT and naive outputs match to numerical precision.
import jax.numpy as jnp
from experiments.jax.week11.hyena_lineage import fftconv, causal_conv1d_naive
B, L, D = 1, 8, 2
u = jnp.arange(B * L * D, dtype=jnp.float32).reshape(B, L, D) / 10.0
k = jnp.zeros((D, L)).at[:, 0].set(1.0).at[:, 1].set(0.5)
bias = jnp.zeros((D,))
y_fft = fftconv(u, k, bias)
y_ref = causal_conv1d_naive(u, k, bias)
print('output shape:', y_fft.shape)
print('max fft/reference difference:', float(jnp.max(jnp.abs(y_fft - y_ref))))
print('first channel:', jnp.round(y_fft[0, :4, 0], 3))
FFT computes circular convolution by default. Padding to 2L makes the circular result agree with linear causal convolution over the first L outputs. Without enough padding, late tokens can wrap around and contaminate early positions.
Guided modify-and-run. Change the kernel so only the last lag is nonzero. A causal implementation should not affect early outputs until that lag becomes available.
Expected observation. Early outputs remain zero for delayed kernels, proving future positions did not leak into the past.
delayed_k = jnp.zeros((D, L)).at[:, 3].set(1.0)
y_delay = fftconv(u, delayed_k, bias)
print('delayed first channel:', y_delay[0, :, 0])
This section makes the notebook intentionally self-contained. The week topic is Hyena, Long Convolution, and Linear Attention Lineage, and the working skill is using padding and frequency-domain multiplication for causal filters. 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/week11/hyena_lineage.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.
The remaining detail to keep in mind is that this notebook is a learning surface, while the tested module is the durable implementation. When the two disagree, trust the tests and inspect the production file. When the notebook adds explanation, use it to form a hypothesis, then confirm that hypothesis against the code path and the printed shape or equivalence check.
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/week11/hyena_lineage.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/week11/ -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 12 moves toward delta-rule and online-learning interpretations, where the model update begins to resemble a learned optimizer.