koopman 0.1.0

Exact custom-gradient ops for the linear Koopman recurrence (PyTorch + JAX)

koopman

Exact, closed-form custom-gradient ops for the linear Koopman recurrence z_{t+1} = K @ z_t. Computes the full state sequence

Z = [z_0, K z_0, K^2 z_0, ..., K^T z_0]

with a single forward scan and an exact adjoint backward — O(T n^2) time and O(T n) memory — instead of taping every per-step matmul through autodiff. Available in a PyTorch flavor (torch.autograd.Function) and a JAX flavor (jax.custom_vjp).

Installation

Install only the backend you need:

pip install koopman[torch]      # PyTorch (CPU or GPU)
pip install koopman[jax-cpu]    # JAX on CPU
pip install koopman[jax-gpu]    # JAX on GPU (CUDA 12)

For a CPU-only PyTorch wheel (e.g. on free CI runners), install torch from the CPU index first, then the package:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install koopman

Requires Python 3.10+.

Usage

PyTorch

import torch
from koopman.torch import linear_powers_sequence

n, T = 16, 1024
K = torch.randn(n, n, dtype=torch.float64, requires_grad=True)
z0 = torch.randn(n, dtype=torch.float64, requires_grad=True)

Z = linear_powers_sequence(K, z0, T)   # [T+1, n]
Z.sum().backward()                     # exact grad_K, grad_z0 via adjoint recurrence

There is also a torch.nn.Module wrapper, LinearPowersSequence(T), and the raw LinearPowersSequenceFn autograd function.

JAX

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from koopman.jax import linear_powers_sequence

n, T = 16, 1024
K = jax.random.normal(jax.random.PRNGKey(0), (n, n))
z0 = jax.random.normal(jax.random.PRNGKey(1), (n,))

Z = linear_powers_sequence(K, z0, T)   # [T+1, n]
loss = lambda K_, z0_: jnp.sum(linear_powers_sequence(K_, z0_, T))
gK, gz0 = jax.grad(loss, argnums=(0, 1))(K, z0)   # jit/vmap-friendly

Why a custom gradient?

Naive autodiff through the Python/lax.scan loop tapes all T matmuls and their activations, costing O(T n^2) memory. The backward here is the closed-form adjoint recurrence:

a_T      = grad_Z[T]
a_t      = grad_Z[t] + Kᵀ @ a_{t+1}        (t = T-1, ..., 0)
grad_K  += outer(a_{t+1}, z_t)
grad_z0  = a_0

It reuses the forward states (O(T n) memory) and avoids materializing the autodiff tape. See benchmarks/ for the comparison against naive autograd and expm-based baselines.

Development

pip install -e ".[dev]"
pytest tests/test_linear_powers_torch.py -v   # PyTorch
JAX_PLATFORMS=cpu pytest tests/test_linear_powers_jax.py -v   # JAX

License

Apache-2.0.