Who contributed to Google on May 9, 2026?

4 developers shipped this update, including copybara-service[bot], TrevorBergeron, daniel-sanche, and chalmerlowe.

What were the notable Google updates?

[Pallas][GPU kernel interpreter] Support simulating multiple GPUs with `shard_map`, [Mosaic GPU] Generate multiple pairwise PTX conversions when possible, and [Mosaic:GPU] Use NCCL API on the host to collective peer addresses.

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 SHIPS MULTI-GPU COLLECTIVE KERNELS AND MOSAIC GPU OPTIMIZATION WAVE

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

4 people shipped this

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

TrevorBergeron @TrevorBergeron 1 cited

daniel-sanche @daniel-sanche 1 cited

chalmerlowe @chalmerlowe 1 cited

JAX's Pallas GPU kernel interpreter now simulates multiple GPUs with shard_map, while Mosaic GPU generates optimized PTX conversions that cut instruction bloat by half [ref:4] [ref:3].

The Pallas GPU kernel interpreter landed three major capabilities overnight: type annotation clarification for callbacks [4], multi-GPU simulation support via shard_map [1], and block/cluster coordinate logging for debugging [5]. On the Mosaic side, engineers shipped an architecture win—generating multiple pairwise PTX conversions in a single inline_asm [2] eliminates the "forest of prmt instructions" LLVM was generating, letting ptxas optimize vector casts dramatically cleaner. The NCCL collective peer address handling [3] now works across processes, unlocking true distributed kernel launches without host-side parameter rendezvous in the symmetrical memory case. Meanwhile, python-genai cut a patch release [6] fixing response_format field naming to snake_case conventions [7], a small but necessary standardization. On the broader cloud SDK front, BigFrames simplified its @udf wrapper object [8], BigTable added client-side metric instrumentation to basic RPCs [9], and db-dtypes dropped Python 3.9 support entirely [10]—a clean break that signals confidence in the 3.10+ baseline.

Action items

→ Test JAX multi-GPU workloads against Pallas shard_map if you're running distributed training google/jax [plan]
→ Upgrade python-genai to 2.0.1 if you're using response_format APIs googleapis/python-genai [immediate]
→ Verify db-dtypes constraint—Python 3.10+ is now mandatory googleapis/google-cloud-python [plan]
→ Monitor BigTable metric instrumentation rollout for performance impact googleapis/google-cloud-python [monitor]

References

[1] [Pallas][GPU kernel interpreter] Support simulating multiple GPUs with `shard_map`. ↗ google/jax
[2] [Mosaic GPU] Generate multiple pairwise PTX conversions when possible ↗ google/jax
[3] [Mosaic:GPU] Use NCCL API on the host to collective peer addresses. ↗ google/jax
[4] [Pallas][GPU kernel interpreter] Clarify type annotations, esp. of Jax/Numpy arrays, of callback arguments. ↗ google/jax
[5] [Pallas][GPU kernel interpreter] Log coordinates of blocks/clusters in the grid. ↗ google/jax
[6] v2.0.1 ↗ googleapis/python-genai
[7] fix: Update response_format field names to snake_case. ↗ googleapis/python-genai
[8] refactor(bigframes): Simplify @udf wrapper object ↗ googleapis/google-cloud-python
[9] feat(bigtable): add client side metric instrumentation to basic rpcs ↗ googleapis/google-cloud-python
[10] fix(db-dtypes): Drop support for Python <= 3.9 ↗ googleapis/google-cloud-python

Quick answers

What shipped in Google on May 9, 2026?: JAX's Pallas GPU kernel interpreter now simulates multiple GPUs with shard_map, while Mosaic GPU generates optimized PTX conversions that cut instruction bloat by half [ref:4] [ref:3]. In total, 38 commits, 35 pull requests, and 1 releases landed.
Who contributed to Google on May 9, 2026?: 4 developers shipped this update, including copybara-service[bot], TrevorBergeron, daniel-sanche, and chalmerlowe.
What were the notable Google updates?: [Pallas][GPU kernel interpreter] Support simulating multiple GPUs with `shard_map`, [Mosaic GPU] Generate multiple pairwise PTX conversions when possible, and [Mosaic:GPU] Use NCCL API on the host to collective peer addresses.

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.