Executable companion — ch01-matrix-exponential.ipynb.

Matrix exponential across array frameworks — NumPy → JAX → PyTorch

Code companion to Chapter 1, section 1.2. The chapter prose is canonical; this notebook is a runnable, output-free companion that puts the three array-framework idioms side by side. All the load-bearing logic lives in the chapter's companion modules (companions/ch01/jax/matrix_exponential.py, companions/ch01/torch/matrix_exponential.py), each pinned by its own pytest suite — this notebook is thin glue that calls them and narrates the contrast.

Run it from the repo root with the companions on the path:

uv run jupyter lab notebooks/ch01-matrix-exponential.ipynb

(The first cell puts the repo root on sys.path so the companions package imports regardless of where Jupyter was launched.)

The computation

The matrix exponential is the power series e^M = sum_{k>=0} M^k / k!. Truncating at order K gives a partial sum S_K. Convergence is fast when the spectral radius of M is small and catastrophically slow when it is large (section 1.2) — which is why production code uses Pade-plus-scaling (scipy.linalg.expm, torch.linalg.matrix_exp) rather than naive truncation.

The partial sums form a recurrence, S_k = S_{k-1} + M^k / k!, and computing the per-order error is a sweep over those partial sums. How each framework expresses that recurrence and that sweep is the whole point of the comparison below.

In [ ]:
# Put the repo root on sys.path so the `companions` namespace imports however
# Jupyter was launched (mirrors pytest's `pythonpath = .`).
import sys
from pathlib import Path
_root = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / "companions").is_dir())
if str(_root) not in sys.path:
    sys.path.insert(0, str(_root))

import numpy as np

import jax
jax.config.update("jax_enable_x64", True)  # JAX: float64 is a *global* config switch

import torch  # PyTorch: precision is set *per tensor* (dtype=torch.float64)

from companions.ch01.jax import matrix_exponential as me_jax
from companions.ch01.torch import matrix_exponential as me_torch

# A 2x2 test matrix with modest spectral radius (the series converges quickly).
A = np.array([[-0.5, 1.0], [-1.0, -0.5]])

1. NumPy — the baseline idiom

Accumulate the running term and running total in an interpreted Python loop. This is the reference both other frameworks are measured against.

In [ ]:
def numpy_truncated_series(M, K):
    """Naive NumPy partial sum S_K via a Python accumulation loop."""
    total = np.eye(M.shape[0])
    term = np.eye(M.shape[0])
    for k in range(1, K + 1):
        term = term @ M / k  # M^k / k!  from  M^{k-1}/(k-1)!
        total = total + term
    return total

numpy_truncated_series(A, 20)

2. JAX — lax.scan + vmap + jit

The partial-sum recurrence becomes a single jax.lax.scan that threads (term, total) as the carry and emits every partial sum in one fused, compiled pass (O(K) instead of NumPy's O(K^2) recompute over a sweep). The per-order error sweep is a jax.vmap. The whole pipeline is jit-compiled; scipy.linalg.expm is the reference.

In [ ]:
# series_partial_sums uses jax.lax.scan; convergence_errors uses jax.vmap.
S_jax = me_jax.series_partial_sums(A, 20)     # shape (21, 2, 2): S_0 .. S_20
errs_jax = me_jax.convergence_errors(A, 20)   # relative Frobenius error vs scipy.expm
S_jax[20], errs_jax[-1]

3. PyTorch — eager loop + batched norm + native matrix_exp

PyTorch is define-by-run and has no scan primitive, so the same recurrence is a plain eager Python loop — the loop is the program. The per-order error is a batched torch.linalg.norm over the stacked partial sums (the counterpart of JAX's vmap), and torch.linalg.matrix_exp is a native reference (JAX borrowed it from SciPy).

In [ ]:
S_torch = me_torch.series_partial_sums(A, 20)     # eager loop -> stacked partial sums
errs_torch = me_torch.convergence_errors(A, 20)   # batched norm vs torch.linalg.matrix_exp
S_torch[20], errs_torch[-1]

The contrast, at a glance

Concern NumPy JAX PyTorch
Partial-sum recurrence Python for loop jax.lax.scan (fused) eager for loop (define-by-run)
Per-order error sweep list comprehension jax.vmap batched torch.linalg.norm
Compilation none (interpreted) jit → XLA eager (opt-in torch.compile)
float64 default global jax_enable_x64 per-tensor dtype
Reference expm scipy.linalg.expm scipy.linalg.expm torch.linalg.matrix_exp (native)

The same lax.scan primitive that fuses this toy recurrence is exactly what powers the S4 / Mamba selective scan in later chapters — which is why the framework's handling of recurrences, not just its array ops, is worth meeting early.

In [ ]:
# All three implement the same math: the truncated series converges to e^A. Each
# companion's pytest suite pins this independently; here we just confirm the JAX and
# PyTorch companions agree with each other (and hence with their native references).
assert np.allclose(errs_jax, errs_torch, atol=1e-12)
assert np.allclose(np.asarray(S_jax[20]), S_torch[20].numpy(), atol=1e-12)
print("JAX and PyTorch agree to 1e-12; both match their native matrix-exponential reference.")