VMEX 0.9.1 — the GPU gradient, the exterior field's grid, and a PyPI install

Optimization gradients now compile on a machine that has a GPU, the coil
examples install from PyPI, and the exterior field sizes its source grid from
the boundary.

Fixed

Optimization gradients could not be compiled on any machine with a GPU.
VMEX deliberately stands the implicit-gradient path down to the CPU on an
accelerator backend, and pinned the host callback to that CPU device — but JAX
requires a pinned device to appear in the enclosing computation's device
assignment, and a jit compiled for cuda:0 does not contain cpu:0. Every
jax.jit-wrapped VmecProblem.jax_value_and_grad died inside JAX's lowering
with ValueError: tuple.index(x): x not in tuple, carrying no VMEX frame and
naming neither the callback nor a device (JAX catches IndexError where
tuple.index raises ValueError, so its own explanatory message never
printed — reported as jax-ml/jax#40722). The pin now applies only within one
platform; the two-accelerator case it exists for keeps it.

Verified on two RTX A4000s, JAX 0.11.1: default placement,
jax.default_device(gpu) and device_scope("gpu") all compile and return the
same value, peak device memory 0.16 GiB, warm value-and-gradient 1.3–1.6 s
against 2.15 s on the same machine's CPU.

The "When the GPU pays off" snippet in the GPU how-to referred to an undefined
runtime and could not be run as written (#157).

Changed

max_fsq_ratio defaults to 1e2, not 1e6. The implicit adjoint assumes
F = 0, so differentiating a trial whose residual is 1e-6 against a 1e-12 deck
carries an O(norm(F)) error and lets a line search walk the design somewhere the
solver cannot resolve; that is what stalled the finite-beta single stage (#361).
At its shipped settings that example now finishes in 900 s and 3.8 GB, driving
its total objective from 31.66 to 4.06 over eight accepted iterations.

The coil examples install from PyPI. ESSOS 0.17 carries uwplasma/ESSOS#58,
so the coils extra pins essos>=0.17 and the
pip install "essos @ git+https://github.com/uwplasma/ESSOS@1b3210ca..."
instruction leaves nine examples, five documentation pages and the examples
README — 100 lines removed. A new all extra installs everything the examples
use. All nine were re-run against plain PyPI essos before the change shipped.

The exterior field sizes its source grid from the boundary.
VmecExtender.from_wout hard-coded nphi=ntheta=32, and the default level
schedule read a per-field-period sampling as a whole-torus count. On the shipped
QA wout (R0/a = 15.9, nfp = 2) the achieved error one minor radius off the
boundary was 2.46e-02 against a requested 1e-6. The sampling now follows the
quadrature's own exp(-2 pi d / h) decay and the schedule carries the nfp
factor: 4.31e-07 at one minor radius and 5.05e-07 at two, at 0.116 s per call
against 0.226 s. A tokamak-like aspect ratio lands on the historical 32 floor,
so its grid does not move; returned fields change only where the old grid was
under-resolved, and explicit nphi, ntheta or levels are honored unchanged.

examples/take_free_boundary_gradients.py certifies against the second adjoint
solver rather than a central difference. On a free boundary the difference
quotient has no usable step — truncation above ~1e-2, solver noise below,
changing sign by 1e-5 — so the 1–2 % it printed read as a gradient error when
the coupled GCROT and edge Schur adjoints agree to 1.6e-04. The fixed-boundary
counterpart, which does have a usable step, certifies at 1.1e-07.

with_near_surface_continuation() now warns that it is broken. At the
default grid it takes 416 s to prepare and returns ~1e-5 in magnitude at every
distance while the direct field falls from 0.52 T at d = 0.25 h to 4e-6 T at
4 h. Where the direct estimate is certified at 3.7e-08 and 1.7e-10 the two
disagree by factors of 4 and 7, so the error is the continuation's. It carries
no error estimate of its own. Use the direct path at a distance its estimate
certifies.

Removed

Four files referenced by nothing in the tree: the QI sheet-current mgrid builder
under tools, a repo census and a stale profile record under benchmarks, and
the pre-commit config, which duplicated tools/preflight.py with a separately
pinned ruff and a 200 kB file cap the tracked tree already exceeds. A repo-wide
audit by reference found nothing else to remove: 9.9 MB across 607 files, every
tool referenced, and benchmarks/baselines/m4/'s 54 records are one per
benchmark cell by design, which is what lets a partial run keep its results.

Coverage

21 of 66 shipped examples were run by no test and nothing said so. Every example
is now parsed by a test, and each must be either named by a test or listed with
a reason. The fixed-boundary gradient example got a test rather than an
exemption: 5.0e-06 at the smoke settings, 1.1e-07 at the shipped ones.

What this release does not fix

Cold compile is still the gap the independent itpplasma/benchmark_vmec corpus
measures. A circular_tokamak solve compiles 152 XLA modules for a 7.6 s cold
run whose warm repeat is 0.16 s, and about 130 of those are single eager JAX ops
each compiling their own module. The persistent cache, on by default, already
halves the cold start across processes (QA_lowres 20.3 s to 10.5 s) and cuts the
single-stage gradient's compile from 49.8 s to 6.0 s, so it is a once-per-machine
cost — but the first run still pays it.

recommended_device still answers gpu for decks where, on two RTX A4000s, the
GPU loses. A full re-run of the CPU-vs-GPU sweep
(benchmarks/gpu_a4000_2026-09-16.json) found no cell where the GPU wins — warm
gain 0.17x to 0.83x across every shipped deck and every point of the synthetic
size scan, including the deck the 2026-07-09 baseline claimed up to 3x for. That
is not a mis-tuned threshold: the crossover does not exist on this hardware and
JAX version. The constant is unchanged, because July's numbers were honestly
measured on other hardware and one machine does not overturn them; the GPU how-to
now carries both so you time your own deck.