test-jax-install.py does not assert the GPU — a silent CPU fallback passes CI in all five repos

Start with scripts/test-jax-install.py in each of the five repositories listed in the issue, comparing the existing jax.devices() call and PRNG-key variant. Run the installation check in the GPU workflow and verify that it fails when JAX uses the CPU and passes when the backend is GPU-backed.

The GPU check that every JAX workflow runs prints the device list and never asserts it. A build whose driver, CUDA install or GPU passthrough had quietly fallen back to the CPU would still pass, and would publish CPU-timed output — which is the one failure GPU CI exists to prevent, since these are GPU-acceleration lectures that time themselves and print their speed gains in the published text.

Verified against main in all five repositories on 2026-09-04.

Where it is

All five copies of scripts/test-jax-install.py are functionally identical — the only differences are the legacy versus current PRNG key call and a trailing newline. None of them contains an assert, a raise or a sys.exit, so the script cannot fail for a wrong-platform reason. Its own comment says the jitted matmul "will likely run on GPU (if available)", which is exactly the uncertainty the check should be removing.

nvidia-smi runs as a separate workflow step, so it proves the host can see a GPU. It does not prove JAX bound to one. Those are different failures: a driver too old for the installed jax[cuda13] wheel gives a perfectly healthy nvidia-smi alongside a CPU-only JAX.

The fix

One line, after the existing jax.devices() call:

assert jax.default_backend() == "gpu", f"JAX fell back to {jax.default_backend()}: {devices}"

jax.default_backend() returns 'gpu' for CUDA and ROCm alike, so it needs no vendor branch. A stricter variant that also names the device is assert any(d.platform == "gpu" for d in jax.devices()).

Why it is worth doing now

Two reasons beyond tidiness.

The first is that the failure is silent by construction. Nothing in the current pipeline distinguishes "the GPU worked" from "JAX ran on the CPU and the lectures published slower numbers", and the timings are part of the rendered output rather than a side effect of it.

The second is that QuantEcon/project-compute (private) is planning a self-hosted GPU runner pilot on hardware we own, where a CPU fallback becomes materially more likely than it is on the rented RunsOn images — the driver becomes ours to maintain, and jax[cuda13] needs an NVIDIA driver of at least 580 and SM 7.5 or newer. The check that ought to catch that regression currently cannot. This surfaced while surveying per-job isolation for that pilot.

Filed here rather than as five issues because the change is the same in all five repositories.