The Wire · Showcase
JAX 0.11.1 lands with export guardrails and a tracing speedup
By RepoJournal · Filed · About Google
JAX's latest release locks the door on expired exports while a core tracing hot path gets a serious speed boost.
JAX v0.11.1 [1] adds a hard error check for deserializing exports older than the compatibility window, closing a class of obscure downstream failures, and introduces a bypass flag for teams that need breathing room. The same release brings jax.numpy.top_k, but the real headline for heavy users is PR #40041 [2], which eliminates quadratic tracing behavior by no longer walking all previously traced equations when building error messages. That fix matters for frameworks like Flax NNX that probe Tracer.sharding and Tracer.device on every operation. Meanwhile, core_map_p is gone [3], with pl.core_map now implemented via mpmd_map_p, and the Pallas team removed MemoryRefBase [4] and transform_ref [5] outright, noting "I decided that a deprecation cycle is not necessary" for the latter. Over on the Python GenAI client, the live connection status enum gains an IDLE state while REQUIRES_ACTION gets deprecated [6]. Cloud Python users should note the bigtable data client now fails V3 mutate_rows entries the server never acknowledged [7], a correctness fix for batch writes, alongside an api_core rest path validation improvement [8].
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
- → Upgrade to JAX v0.11.1 and verify your export deserialization workflows against the new compatibility check google/jax [immediate]
- → Test any code using core_map_p or the removed Pallas helpers against the new implementations google/jax [immediate]
- → Review bigtable batch mutation handling for unacknowledged rows after upgrading google-cloud-python googleapis/google-cloud-python [plan]
- → Update live connection status enums to use IDLE instead of REQUIRES_ACTION googleapis/python-genai [plan]
References
- [1] JAX v0.11.1 ↗ google/jax
- [2] PR #40041: Avoid quadratic tracing behaviour google/jax
- [3] Removed `core_map_p` and everything which depended on it google/jax
- [4] [pallas] Removed `MemoryRefBase` google/jax
- [5] [pallas:mgpu] Removed `transform_ref` google/jax
- [6] feat: Add IDLE state to live connection status enum and mark REQUIRES_ACTION as deprecated. googleapis/python-genai
- [7] fix(bigtable): data client should acknowledge all mutations in batch ↗ googleapis/google-cloud-python
- [8] fix(api_core): improve rest path validation ↗ googleapis/google-cloud-python