RepoJournal
Google

Google

JAX, the GenAI SDK, and the Cloud libs - Google's open source layer

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.

Pick a date

Topics: AI / ML Python Full archive →

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.

Action items

References

  1. [1] [vjp3][mutable-arrays] fix scan transpose bugs with refs ↗ google/jax
  2. [2] [JAX] Enable CHLO lowering for cumulative sums on GPU again. ↗ google/jax
  3. [3] Update XLA dependency to use revision http://github.com/openxla/xla/commit/42d2e6ac6c715184ab7a64855f5162f03c561b29 google/jax

Quick answers

What shipped in Google on July 5, 2026?
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. In total, 3 commits and 2 pull requests landed.
Who contributed to Google on July 5, 2026?
2 developers shipped this update, including mattjj and copybara-service[bot].
What were the notable Google updates?
[vjp3][mutable-arrays] fix scan transpose bugs with refs, [JAX] Enable CHLO lowering for cumulative sums on GPU again, and Update XLA dependency to use revision http://github.com/openxla/xla/commit/42d2e6ac6c715184ab7a64855f5162f03c561b29.

More from Google

Daily updates, in your inbox

Follow Google

Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.

or

One email a day. Unsubscribe in one click. Read a past issue →

Elsewhere on the wire

Want every project, not just this one?

We use privacy-friendly analytics (Google Analytics, IP-anonymized) to see which pages help readers. No ads, and we never sell your data. See our Privacy Policy.