JAX Primitives Through Dynamical Systems

Motivation

The standard route into state space models begins with lax.scan tutorials that implement toy counters or fibonacci sequences. We take a different path: implement dynamical systems you have solved before, and let the JAX primitives emerge as the natural computational substrate.

This chapter builds four pieces of infrastructure that recur throughout the curriculum:

1. A lax.scan-based time stepper (used in every SSM implementation),
2. Automatic Jacobian computation via jax.jacfwd (stability analysis),
3. QR-based Lyapunov exponents via lax.scan (training diagnostics, Week 18),
4. The associative scan operator for parallel prefix computation (S5, Mamba).

The damped oscillator as scan prototype

Continuous system

Consider the damped harmonic oscillator:

where $\omega$ is the natural frequency and $\gamma \geq 0$ is the damping coefficient. In state-space form with state $\statevec = [x, v]^\top$:

The eigenvalues of $\statemat$ are $\lambda = -\gamma/2 \pm \sqrt{(\gamma/2)^2 - \omega^2}$. For underdamped motion ($\gamma < 2\omega$), these are complex with negative real part — the system oscillates and decays. This eigenvalue structure is identical to what we will see in trained S4D layers (Week 5).

Euler discretization maps to lax.scan

Forward Euler with step size $\Delta t$:

This is a recurrence of the form $\statevec_{n+1} = f(\statevec_n)$ — exactly the pattern that lax.scan computes. The scan signature (carry, x) -> (carry, y) maps as:

Energy analysis

The energy $E = \frac{1}{2}(v^2 + \omega^2 x^2)$ satisfies $\dot{E} = -\gamma v^2 \leq 0$. For $\gamma > 0$, energy decays monotonically in continuous time.

Gradient sensitivity

JAX’s jax.grad computes $\partial x_{\text{final}} / \partial \gamma$ — the sensitivity of the endpoint to the damping coefficient. This works through the entire lax.scan via reverse-mode automatic differentiation.

fori_loop and jit: alternatives to scan

lax.fori_loop compiles the same recurrence but returns only the final state — no trajectory stored. This is the inference pattern: when generating tokens, you need $\statevec_L$, not the intermediate $\statevec_1, \ldots, \statevec_{L-1}$.

jax.jit wraps the entire scan into a single XLA computation, eliminating Python interpreter overhead between steps. The first call traces and compiles; subsequent calls dispatch directly to the compiled kernel.

Coupled oscillators and the Jacobian

Ring topology with nearest-neighbor coupling

$N$ oscillators on a ring, coupled via discrete Laplacian:

The state vector is $\statevec = [x_1, \ldots, x_N, v_1, \ldots, v_N]^\top \in \R^{2N}$.

Jacobian as transition matrix

The Jacobian of one Euler step with respect to the state, $\jacobian = \partial f / \partial \statevec$, is the matrix of all first-order partial derivatives — it captures how a small perturbation in $\statevec$ propagates through one time step. This makes $\jacobian$ the discrete-time transition matrix: it describes how the state evolves locally.

For linear systems, $\jacobian$ is constant (independent of $\statevec$) — the same transition applies everywhere. For nonlinear systems, $\jacobian$ varies along the trajectory. In the SSM context, this distinction maps to Linear Time-Invariant (LTI) models (S4, S4D) where parameters are fixed, versus selective models (Mamba) where parameters depend on the input. ”Selective” means the model can choose what to remember based on input content, not just position.

We compute $\jacobian$ via jax.jacfwd(one_step)(state0), which returns the full $2N \times 2N$ Jacobian matrix.

Batching with vmap

jax.vmap vectorizes a function over a batch dimension without writing explicit loops. It transforms a function $f: \R^n \to \R^m$ into $\hat{f}: \R^{B \times n} \to \R^{B \times m}$ that processes $B$ inputs simultaneously on GPU.

Two demonstrations:

1. Ensemble simulation: vmap over initial conditions produces $B$ oscillator trajectories in parallel. This is the same pattern used to process a minibatch of sequences through an SSM layer.
2. Batched Jacobians: vmap(jacfwd) computes the transition matrix at $B$ different states simultaneously. In Week 18, this enables Lyapunov analysis across minibatches of sequences.

Lyapunov exponents via QR decomposition

The numerical catastrophe

Given a sequence of transition matrices $\jacobian_0, \jacobian_1, \ldots, \jacobian_{L-1}$, the cumulative effect is the product $M_L = \prod_{t=0}^{L-1} \jacobian_t$. Computing $M_L$ directly is numerically catastrophic: if any eigenvalue has $\abs{\lambda} > 1$, components grow exponentially ($\sim \abs{\lambda}^L$), causing overflow. If $\abs{\lambda} < 1$, components decay to zero (underflow). With $L = 10{,}000$ steps, even $\abs{\lambda} = 1.001$ grows to $e^{10} \approx 22{,}000$.

The QR algorithm

The solution is to periodically orthogonalize via QR decomposition. [Official] Benettin et al. Benettin et al. (1980) introduced the algorithm; see Anonymous (2025) for a recent application in the SSM context. At each step:

1. Propagate the perturbation matrix: $M = \jacobian_t Q_{t-1}$
2. Decompose: $Q_t R_t = M$ (QR factorization)
3. Accumulate: $\text{log\_sum} \mathrel{+}= \log \abs{\diag(R_t)}$

The Lyapunov exponents are the time-averaged growth rates:

The computation fits naturally into lax.scan. Its carry stores the state, the orthogonal basis, and the running log sum:

This is the same carry-and-update pattern used by SSM recurrences.

Parallel prefix computation

Warm-up: cumulative sum

The simplest associative_scan example: cumulative sum with the operator $f(a, b) = a + b$. Addition is trivially associative, so:

def cumsum_associative(x):
    return jax.lax.associative_scan(jnp.add, x)

This produces identical results to jnp.cumsum but makes the parallel prefix structure explicit. The Blelloch algorithm Blelloch (1990) computes all $L$ prefix sums in $\bigO{\log \seqlen}$ parallel depth — a general technique for reducing an associative operator over a sequence in $\bigO{\log L}$ parallel steps using $\bigO{L}$ total work. Same FLOPs as a sequential scan, but vastly more parallelism on GPU.

The EMA recurrence

The exponential moving average (EMA) is a running weighted average that gives more weight to recent values. With smoothing parameter $\alpha \in (0, 1)$:

When $\alpha$ is small (e.g., 0.1), the average is “slow” and remembers far into the past. When $\alpha$ is large (e.g., 0.9), it tracks the input closely. This is a first-order linear recurrence — the simplest possible SSM with scalar state, $\discA = (1-\alpha)$, and $\discB = \alpha$.

Sequential computation via lax.scan requires $\bigO{\seqlen}$ sequential steps. Can we do better?

The associative operator

Define the pair representation $(a_t, b_t) = (1-\alpha, \, \alpha \, x_t)$, where $a_t$ is the decay factor and $b_t$ is the input contribution. We need a binary operator — a function that combines two such pairs into one, representing the combined effect of two consecutive time steps. The key requirement is associativity: $(p \scanop q) \scanop r = p \scanop (q \scanop r)$, which allows arbitrary grouping and therefore parallel evaluation. The operator:

Associativity is what enables jax.lax.associative_scan to compute the prefix sums in $\bigO{\log \seqlen}$ parallel depth using Blelloch’s algorithm.

Reflections

What transfers directly from the PhD

Three pieces of infrastructure from vortex dynamics research reappear with zero modification:

1. QR-based Lyapunov computation. Same algorithm, same numerical considerations (transient removal, convergence diagnostics). The only adaptation: SSM dynamics are input-dependent, so the spectrum is a statistic over sequences rather than a property of the system.
2. Discretization awareness. The intuition that “choice of integrator determines long-time accuracy” maps directly to “choice of SSM discretization determines training stability.” The stability region analysis from numerical methods applies unchanged.
3. Eigenvalue monitoring. Tracking whether eigenvalues stay in the left half-plane (continuous) or inside the unit disk (discrete) is the same diagnostic in both domains.

What requires adaptation

1. No natural period. Floquet theory assumes periodicity; SSM dynamics depend on the input sequence. The correct analogue is Lyapunov exponents (infinite-time average) rather than Floquet multipliers (one-period product).
2. High dimension. The vortex quartet has 4–8 degrees of freedom. SSMs have $\statedim = 16$–$256$ per layer, with $L = 1024$–$131072$ time steps. The QR algorithm scales well ($\bigO{d^2 L}$), but visualization strategies change.
3. Learned dynamics. In the dissertation, $\statemat$ is given by physics. In SSMs, $\statemat$ is learned — meaning the stability properties change during training. This opens the question of whether stability-aware training (constraining eigenvalues) improves optimization.

Forward map

• Week 2: The training loop harness built in Flax will use lax.scan for the SSM layer and optax for optimization.
• Weeks 4–6: S4, S4D, and S5 use the exact scan patterns from this chapter, with HiPPO-initialized $\statemat$ matrices instead of our hand-crafted oscillator systems.
• Week 7: Mamba makes $\inputmat$, $\outputmat$, and $\stepsize$ input-dependent, breaking the convolutional view but preserving the parallel scan.
• Week 18: The Lyapunov infrastructure from Exercise 3 is applied to trained SSMs to characterize gradient flow stability.
• Week 20: The discretization analysis from the DynConnect boxes is formalized with stability region plots in Julia.