Who contributed to Google on May 23, 2026?

1 developer shipped this update, including copybara-service[bot].

What were the notable Google updates?

PR #35699: Strip Mosaic kernel debug info and exclude autotune cache dir from cache keys, Add tests for uneven axis index groups in collectives, and Remove stale version tests.

Google

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

Pick a date

Topics: AI / ML Python Full archive →

The Wire · Showcase

JAX FIXES CACHE KEY BUGS BREAKING DISTRIBUTED TRAINING ACROSS ENVIRONMENTS

By RepoJournal · Filed 06:03 UTC on May 23, 2026 · About Google

1 person shipped this

copybara-service[bot] @copybara-service[bot] 3 cited

JAX's compilation cache is now shipping with environment-agnostic keys, fixing a critical issue where autotune paths were poisoning builds across machines.

The Mosaic kernel cache fix [1] strips debug info from TPU custom calls that were bloating cache keys and causing unnecessary recompilation. This lands alongside a collective operations fix [2] that corrects a regression in psum and pmean when handling uneven axis index groups with constants. The real impact: distributed training jobs that span multiple hardware configs will stop thrashing the autotune cache on every run. JAX also bumped minimum version floors [3], dropping support for jaxlib_extension versions below 452 and libtpu below 0.0.41. Breaking change in Pallas: `pl.program_id` no longer works inside `plgpu.kernel` [4], but the migration path is clean - swap to `lax.axis_index` and `lax.axis_size`. On the Python SDK side, google-cloud-python dropped Python 3.7-3.9 support [5] and fixed a pubsub generation issue [6] that was blocking client library builds. The GenAI SDK added reinforcement tuning support [7], moving one step closer to a complete fine-tuning API.

Action items

→ Pin JAX to the latest commit with cache key fix before scaling distributed jobs jax-ml/jax [plan]
→ Update google-cloud-python past version supporting Python 3.7-3.9 removal googleapis/google-cloud-python [plan]
→ If using Pallas multi-GPU kernels, migrate pl.program_id to lax.axis_index jax-ml/jax [monitor]

References

[1] PR #35699: Strip Mosaic kernel debug info and exclude autotune cache dir from cache keys google/jax
[2] Add tests for uneven axis index groups in collectives. ↗ google/jax
[3] Remove stale version tests. google/jax
[4] [pallas:mgpu] `plgpu.kernel` is now implemented in terms of `pl.kernel` ↗ google/jax
[5] chore: remove references to Python 3.7, 3.8, and 3.9 (#17227) googleapis/google-cloud-python
[6] fix: resolve issue with pubsub generation (#17219) googleapis/google-cloud-python
[7] feat: Support Reinforcement Tuning in GenAI SDK ↗ googleapis/python-genai

Quick answers

What shipped in Google on May 23, 2026?: JAX's compilation cache is now shipping with environment-agnostic keys, fixing a critical issue where autotune paths were poisoning builds across machines. In total, 53 commits and 50 pull requests landed.
Who contributed to Google on May 23, 2026?: 1 developer shipped this update, including copybara-service[bot].
What were the notable Google updates?: PR #35699: Strip Mosaic kernel debug info and exclude autotune cache dir from cache keys, Add tests for uneven axis index groups in collectives, and Remove stale version tests.

JAX FIXES REMAT3 CONSTANT HANDLING CRASH, GOOGLE CLOUD PYTHON CENTRALIZES REST TRANSCODING

Matthew Johnson patched a critical remat3 regression where closed-over constants were silently dropped, causing immediate crashes in production rematerialization pipelines [ref:1].

ai 50 shipped 1-min read

Google 1 day ago

JAX TIGHTENS XLA OPTIMIZATION CONTROLS, GOOGLE CLOUD PYTHON SETS UP FEATURE GATES

JAX makes XLA optimization levels human-readable while Google Cloud Python client libraries prepare a unified feature gating system for observability.

ai 54 shipped 1-min read

Google 2 days ago

GENAI AUDIO TRANSCRIPTION SHIPS, PROTO-PLUS FIXES RACE CONDITIONS

Google's AI SDK gains native audio processing while the core Python cloud libraries ship critical stability fixes across the stack.

ai 60 shipped 1-min read

Google 3 days ago

JAX CUTS ASYNC COLLECTIVE PRIMITIVES, GENAI SDK ADDS LIVE API VOCABULARY

JAX is simplifying its async collective operations while the Python GenAI SDK ships custom vocabulary support for the Live API, marking a shift toward leaner abstractions and richer real-time capabilities.

ai 100 shipped 1-min read

Elsewhere on the wire

AI Agents about 9 hours ago

CLAUDE OPUS 5 LANDS ACROSS THE STACK

The newest Anthropic model is now live in langchain, Cline, and llama-index, with native support for extended reasoning and 1M context windows.

ai-agents 28 shipped 1-min read

Local LLMs about 9 hours ago

OLLAMA LANDS LAGUNA SUPPORT AND CRUSHES MEMORY LEAKS WHILE SGLANG HITS V0.5.16 WITH CONFIDENCE-DRIVEN SPECULATIVE DECODING

Ollama shipped three critical performance and reliability fixes for Metal residency and concurrent access patterns, while SGL-Lang released 0.5.16 with a new speculative algorithm hitting 383.7 tok/s on DeepSeek-V4.