jax 0.11.1: XLA:CPU dynamic-update-slice-in-loop regression (why jax is pinned to 0.11.0)

This issue is the public record of why every jax install in this repo is pinned to 0.11.0, what the underlying jax/XLA regression actually is, and what has to be true before the pins can lift. It exists so the upstream report has a public reference for where the bug was found and how it was diagnosed.

The regression, in one paragraph

Since jax/jaxlib 0.11.1 (PyPI 2026-08-17), a dynamic-update-slice that writes a small slice (< 256 bytes) into a large array inside a loop body costs O(whole buffer) per iteration on the CPU backend instead of O(update). In the standard lax.fori_loop / lax.scan accumulation idiom the buffer length equals the trip count, so runtime becomes quadratic in the loop length — ~4x per doubling of n. This lecture's numpy_vs_numba_vs_jax.md runs exactly that idiom at n = 10,000,000: ~0.06 s under 0.11.0, extrapolated ~18 hours under 0.11.1, so every executing build dies on CellTimeoutError at the 600 s myst-nb limit.

How it surfaced and what landed

All six workflows in this repo and all nine translation workflows now pin 0.11.0. Pinning jax alone pins jaxlib too: jax==0.11.0 requires jaxlib<=0.11.0,>=0.11.0 on the base requirement and on every accelerator extra, and the pairing is enforced again at import in both directions.

Root cause, condensed

The full report with the reproduction script and measurement tables is being filed upstream at jax-ml/jax (link to follow in a comment). The short version, all measured:

• Bisected to a single nightly build: 0.11.1.dev20260725 good (n=200,000 fori 0.0014 s), 0.11.1.dev20260726 bad (13.24 s) — ~10,000x apart, monotonic on both sides, reproduced on linux/x86_64 and linux/aarch64.
• It is jaxlib, not jax: mixing nightly packages, runtime tracks the jaxlib version only, and the optimized HLO is byte-identical between fast and slow pairings. The boundary corresponds to the XLA roll 6b5d5254... -> 88e9a7db..., a window of exactly 10 commits, 4 tagged [XLA:CPU].
• The trigger is the DUS, not the loop: fori_loop, scan and while_loop with scalar bodies are unaffected on 0.11.1; the identical scan with its stacked ys output discarded is unaffected; only bodies writing a small slice into a large buffer regress. Cost is O(buffer x trips) with a sharp cliff at exactly 256 bytes of update width (holds across float16/32/64 — a byte threshold, not an element count).
• No flag-level workaround exists: --xla_cpu_use_fusion_emitters was deprecated inside the same XLA window, and the surviving related flags measurably change nothing. The only mitigations are pinning, discarding stacked outputs, or batching writes to >= 256 bytes per iteration.

Unpin condition

A jax release must ship whose jitted lax.fori_loop at n=400,000 on CPU completes in ~single-digit seconds rather than ~100 s (the two sides differ by ~65x, so the gate is unambiguous; a validated one-command container test exists). As of 2026-08-20 no such release exists — 0.11.1 is latest, and the nightly 0.11.2.dev20260819 still carries the regression at ~96% of 0.11.1's runtime, so a hypothetical 0.11.2 cut from current main would ship broken. When the upstream fix lands, unpin the six workflows here, the nine translation workflows, and the lecture's install cell together.

Diagnosed with assistance from Anthropic's Claude models (Opus 5 and Fable 5).

从 cache.yml、ci.yml、publish.yml、三个执行工作流、翻译工作流以及讲义 lecture's numpy_vs_numba_vs_jax.md 的安装单元开始。在更改 pin 之前，先验证更新的 JAX 版本是否通过针对 n=400,000 所规定的 CPU 计时门槛。修复后，同时提升所有列出的 pin，并确认讲义和工作流执行都能在现有时间限制内完成。