The Wire · Showcase
JAX FIXES CRITICAL SCAN TRANSPOSE BUG THAT BROKE REF ACCUMULATORS
By RepoJournal · Filed · About Google
A foundational bug in JAX's scan transpose logic is fixed after the code assumed accumulators would follow a neat type-sorted pattern they never actually did.
The scan transpose implementation was making a dangerous assumption about how ref accumulators would be arranged in memory [1]. The code expected them to be sorted and contiguous by type, but real-world usage scattered them arbitrarily across the accumulator space. This created silent correctness issues that the new canonicalization logic now handles properly, complete with systematic test coverage to prevent regression [1]. In parallel, JAX is restoring cumulative sum performance on GPU by re-enabling CHLO lowering after a recent disable [2]. XLA dependencies were also bumped to the latest revision to align with upstream improvements [3].
One email a day. Unsubscribe in one click.
Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.
One email a day. Unsubscribe in one click. Read a past issue →
Action items
- → Pull the latest JAX commit with ref accumulator fix [ref:1] if you use scan with mutable refs google/jax [immediate]
- → Test cumulative sum operations on GPU to verify performance is back to baseline google/jax [monitor]
References
- [1] [vjp3][mutable-arrays] fix scan transpose bugs with refs ↗ google/jax
- [2] [JAX] Enable CHLO lowering for cumulative sums on GPU again. ↗ google/jax
- [3] Update XLA dependency to use revision http://github.com/openxla/xla/commit/42d2e6ac6c715184ab7a64855f5162f03c561b29 google/jax