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
| Repository |
Variant |
lecture-jax |
jax.random.PRNGKey(0) |
lecture-python.myst |
jax.random.key(0) |
lecture-python-programming |
jax.random.PRNGKey(0) |
lecture-stats |
jax.random.PRNGKey(0) |
iuj_feb_2026 |
jax.random.PRNGKey(0) |
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.
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
mainin all five repositories on 2026-09-04.Where it is
lecture-jaxjax.random.PRNGKey(0)lecture-python.mystjax.random.key(0)lecture-python-programmingjax.random.PRNGKey(0)lecture-statsjax.random.PRNGKey(0)iuj_feb_2026jax.random.PRNGKey(0)All five copies of
scripts/test-jax-install.pyare functionally identical — the only differences are the legacy versus current PRNG key call and a trailing newline. None of them contains anassert, araiseor asys.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-smiruns 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 installedjax[cuda13]wheel gives a perfectly healthynvidia-smialongside a CPU-only JAX.The fix
One line, after the existing
jax.devices()call:jax.default_backend()returns'gpu'for CUDA and ROCm alike, so it needs no vendor branch. A stricter variant that also names the device isassert 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, andjax[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.