From fc0b03d3a111bab8ee8c75479a32fb8df75bb0ea Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Wed, 19 Aug 2026 12:16:52 +1000 Subject: [PATCH] Pin jax to 0.11.0 in CI -- jax 0.11.1 hangs the fori_loop cells on CPU jax 0.11.1 (2026-08-17) regresses lax.fori_loop on the CPU path. numpy_vs_numba_vs_jax's jax cells jit with device=cpu, so the regression bites even on this repo's GPU runners; the unpinned 'pip install -U jax[cuda13]' would resolve the broken release at the next Monday cache rebuild (2026-08-24). Executed in 8.18s under 0.11.0 on 2026-08-17; the fr/fa siblings measured >1800s under 0.11.1. Full evidence and the family-wide exposure map are recorded in QuantEcon/workspace-lectures#49. Co-Authored-By: Claude Fable 5 --- .github/workflows/cache.yml | 8 +++++++- .github/workflows/ci.yml | 8 +++++++- .github/workflows/publish.yml | 8 +++++++- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/.github/workflows/cache.yml b/.github/workflows/cache.yml index 2baed89a..981325bf 100644 --- a/.github/workflows/cache.yml +++ b/.github/workflows/cache.yml @@ -21,7 +21,13 @@ jobs: - name: Install JAX and Numpyro shell: bash -l {0} run: | - pip install -U "jax[cuda13]" + # Pinned: jax 0.11.1 (2026-08-17) hangs numpy_vs_numba_vs_jax's + # lax.fori_loop cells -- they jit with device=cpu, so the CPU-path + # regression bites even on these GPU runners. Context is recorded in + # QuantEcon/workspace-lectures#49. With jax already present, the + # lecture's unpinned `!pip install quantecon jax` cell is satisfied + # and does not upgrade. + pip install "jax[cuda13]==0.11.0" pip install numpyro python scripts/test-jax-install.py - name: Check nvidia drivers diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b36ab89d..e5ed1e84 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,7 +26,13 @@ jobs: - name: Install JAX and Numpyro shell: bash -l {0} run: | - pip install -U "jax[cuda13]" + # Pinned: jax 0.11.1 (2026-08-17) hangs numpy_vs_numba_vs_jax's + # lax.fori_loop cells -- they jit with device=cpu, so the CPU-path + # regression bites even on these GPU runners. Context is recorded in + # QuantEcon/workspace-lectures#49. With jax already present, the + # lecture's unpinned `!pip install quantecon jax` cell is satisfied + # and does not upgrade. + pip install "jax[cuda13]==0.11.0" pip install numpyro python scripts/test-jax-install.py - name: Install latex dependencies diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index e61eb6d6..f832593d 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -38,7 +38,13 @@ jobs: - name: Install JAX and Numpyro shell: bash -l {0} run: | - pip install -U "jax[cuda13]" + # Pinned: jax 0.11.1 (2026-08-17) hangs numpy_vs_numba_vs_jax's + # lax.fori_loop cells -- they jit with device=cpu, so the CPU-path + # regression bites even on these GPU runners. Context is recorded in + # QuantEcon/workspace-lectures#49. With jax already present, the + # lecture's unpinned `!pip install quantecon jax` cell is satisfied + # and does not upgrade. + pip install "jax[cuda13]==0.11.0" pip install numpyro python scripts/test-jax-install.py - name: Check nvidia drivers