Pin jax in the numpy_vs_numba_vs_jax install cell - #621
Conversation
The workflow pin from #620 protects the build but not the reader. In CI the cell is a no-op because jax 0.11.0 is already installed, so pip treats the unpinned requirement as satisfied. In Colab nothing pre-satisfies it: a reader executing this notebook today resolves jax 0.11.1, then reaches the CPU-pinned lax.fori_loop and lax.scan cells at n = 10,000,000. jax 0.11.1 regresses XLA:CPU execution quadratically -- measured on linux x86_64, doubling n multiplies runtime by ~4.1 (fori) and ~4.0 (scan). At n=400,000 it is 105s against 0.005s under 0.11.0; extrapolated to the lecture's n=10,000,000 that is roughly 18 hours, which a reader experiences as a hang. Pinned to ==0.11.0 rather than !=0.11.1 deliberately: the regression is still present on jax main (nightly 0.11.2.dev20260819 measured at 96% of 0.11.1's time), so an exclusion would admit a likely-broken 0.11.2 without warning. This matches the exact version the six workflows install, so readers now run what CI runs. The other two jax cells in this repo are untouched: jax_intro.md loops to n=20 and autodiff.md has no lax loop, so neither reaches the regime. Lifting this pin is tracked alongside the workflow pins in QuantEcon/workspace-lectures#49. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Pins the JAX version installed by the numpy_vs_numba_vs_jax lecture’s notebook cell so readers (e.g., in Colab) run the same JAX version as CI, avoiding performance regressions from newer releases.
Changes:
- Update the lecture’s
pip installnotebook cell to installjax==0.11.0explicitly.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
✅ Translation sync completed (fa)Target repo: QuantEcon/lecture-python-programming.fa
|
✅ Translation sync completed (zh-cn)Target repo: QuantEcon/lecture-python-programming.zh-cn
|
✅ Translation sync completed (fr)Target repo: QuantEcon/lecture-python-programming.fr
|
Pins the jax install cell in
numpy_vs_numba_vs_jax.mdto0.11.0, the version the six workflows already install.Why this is not covered by #620
#620 pinned the build. It cannot protect the reader, and the two paths differ in exactly the way that matters:
!pip install quantecon jaxdoesn = 10_000_000fori_loop / scanThe published notebook is live now —
python-programming.quantecon.org/_notebooks/numpy_vs_numba_vs_jax.ipynb(HTTP 200, 28,740 B) and the generated mirror both still ship the unpinned cell next ton = 10_000_000and the CPU-pinned loops.Why
==0.11.0and not!=0.11.1An exclusion looks tidier and self-heals when a fix ships, but it is the wrong choice today: the regression is still present on jax main. The nightly
0.11.2.dev20260819was measured in a container at 96% of 0.11.1's runtime (fori 100k/200k/400k = 6.21 / 25.34 / 100.59 s, against 0.005 s flat under 0.11.0), so a released 0.11.2 would very likely carry the bug and!=0.11.1would admit it silently.Pinning to the exact version CI installs also means readers now execute what the build executes, which is the property the lecture wants anyway — it is a performance comparison, and a reader on a different jax is not measuring the same thing.
Measurement behind the claim
jax 0.11.1 regresses XLA:CPU execution quadratically. Same container, same script, only the jax version differing (second, post-compilation run, seconds):
Doubling ratios are 4.12 / 4.10 for fori and 3.98 / 3.96 for scan — O(n²) against an ideal 4.0. A stack sample of the stalled process sits in
xla::cpu::ThunkExecutor::ExecuteSequential, so compilation has already finished and this is runtime execution. Reproduced on linux x86_64, linux aarch64 and macOS arm64.Scope
Only this lecture changes.
jax_intro.mdloops to n=20 andautodiff.mdhas no lax loop, so neither reaches the regime and both keep their unpinned cells.The same cell exists in the three translations at
numpy_vs_numba_vs_jax.md:64. Translation sync is.md-based, so it should carry this across rather than needing three hand-mirrored PRs — worth confirming on the next sync run.Lifting this pin belongs with the workflow pins, tracked on QuantEcon/workspace-lectures#49. Nothing upstream to wait on yet: there is no jax release newer than 0.11.1, and no upstream issue describing this regression exists.