JAX performance notes

A short, practical guide to compiling and calling lyapax efficiently. None of this changes results - it’s about not paying more XLA compilation cost than necessary.

What’s traced vs. what’s Python

lyapunov_spectrum, lyapunov_spectrum_dde, and sweep_lyapunov_spectrum are ordinary Python functions, not jax.jit-decorated ones. Internally they use n_steps, renorm_every, k, and t_transient as plain Python values to size jax.lax.scan lengths and pre-allocate the tangent matrix ((d, k)) before entering traced code - none of that is itself run through XLA tracing at the top level. state0, dt, params, and (for DDEs) buf0 are the values that flow into the traced/compiled inner jax.jvp/jax.vmap/jax.lax.scan machinery.

Wrapping a call in jax.jit

If you wrap lyapunov_spectrum/lyapunov_spectrum_dde in your own jax.jit (e.g. to fuse it into a larger pipeline), mark the shape/control-flow-determining arguments as static:

import functools
import jax
from lyapax import lyapunov_spectrum

jitted = jax.jit(
    lyapunov_spectrum,
    static_argnames=["n_steps", "k", "renorm_every"],
)

n_steps, k, and renorm_every determine jax.lax.scan lengths and array shapes (the (d, k) tangent matrix, the n_steps // renorm_every loop trip count) - changing any of them triggers a full XLA recompile under jax.jit, static or not. t_transient is used to compute a step count too (renorm_every * round(t_transient / dt / renorm_every)), so it is effectively shape-determining as well; treat it as static (or fixed) in code paths you plan to call repeatedly.

In practice you often don’t need to jax.jit the call yourself: the jax.jvp/jax.vmap/jax.lax.scan calls inside lyapunov_spectrum already get compiled by JAX on first use with a given set of static shapes and are cached across subsequent calls with the same shapes - see 07_speed_and_accuracy.py and Benchmarks, which report wall time after this warmup compilation.

Compilation cost vs. runtime

The first call for a given (d, k, n_steps, renorm_every, dt, integrator) combination pays an XLA tracing/compilation cost on top of the actual integration; later calls with the same static shapes reuse the compiled executable and only pay runtime. When benchmarking or reporting timings, always report compile-time (first call) and steady-state runtime separately, as Benchmarks does - a single combined number overstates per-call cost for any workload that calls lyapunov_spectrum more than once with the same shapes, and understates it for a true one-shot call.

Avoiding accidental recompilation

• Don’t vary n_steps/k/renorm_every across a Python loop if you can help it - each distinct value triggers a new compile. If you need to sweep a value (a parameter, an initial condition), sweep it as data instead: see sweep_lyapunov_spectrum and 11_vmap_parameter_sweep.py, which batches a whole grid into one jax.vmap-compiled call rather than one Python-level call (and one compile) per grid point.

• Reuse dt, n_steps, k, renorm_every across repeated calls (e.g. inside an optimization loop or a notebook cell you re-run) so the compiled executable is reused instead of rebuilt.

• Don’t put lyapunov_spectrum/lyapunov_spectrum_dde calls inside a Python loop with changing static shapes expecting jax.jit caching to save you - caching is keyed on the static shapes, and a different n_steps or k is a cache miss, not a fast path.