$ the-wire · showcase
JAX handles non-unimplemented layout errors, bumps XLA
By RepoJournal · Filed · About Google · Composed from the cited sources · methodology
JAX now gracefully handles errors other than "unimplemented" when reading layouts, closing a crash path for arrays used after deletion or donation.
JAX changed how it responds when reading layouts fails. Previously, an error other than "unimplemented" during layout reads could surface as a crash; the new behavior handles those errors gracefully and avoids crashes in tests that try to use arrays after their deletion or donation, where a runtime error was being raised [1]. A dedicated test was added to exercise this behavior [1]. The same change landed through the JAX PR pipeline [2].
The other JAX change this cycle is a dependency move: the XLA dependency was updated to a new revision [3]. That means builds pulling JAX now track the newer XLA commit, so anyone pinning or vendoring XLA should align with the revision JAX expects [3].
Neither change is disruptive on its own. The layout error handling is defensive and comes with a test; the XLA bump is a routine revision update that will matter most to consumers building against a fixed XLA SHA [3].
Action items
- → Run JAX array-after-donation tests against the new layout error handling google/jax [plan]
- → Align vendored XLA revision with JAX's updated dependency google/jax [monitor]
References
- [1] Gracefully handle errors other than "unimplemented" when reading layouts. ↗ google/jax
- [2] Gracefully handle errors other than "unimplemented" when reading layouts. ↗ google/jax
- [3] Update XLA dependency to use revision http://github.com/openxla/xla/commit/b7fff55c413a9057a0e0c76503fe2f41598d05d0 ↗ google/jax