Who contributed to Google on May 31, 2026?

2 developers shipped this update, including mattjj and copybara-service[bot].

What were the notable Google updates?

restore triangular_solve batching rule, improve sharding rule and Support while_p in GPU Interpret Mode.

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 TRIANGULAR_SOLVE MEMORY BLOWUP, EXPANDS GPU INTERPRET MODE

By RepoJournal · Filed 06:02 UTC on May 31, 2026 · About Google

2 people shipped this

mattjj @mattjj 1 cited

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

A partial revert of yesterday's batching rule is rolling back OOMs in triangular_solve while keeping the performance gains.

JAX shipped a fix [1] that reverts the problematic broadcast insertion from the previous batching rule but preserves the optimization by improving the sharding rule instead. That earlier PR was causing out-of-memory errors in certain batching cases, so this lands the best of both worlds: the batching rule stays lean while the sharding rule now accepts additional batch dimensions. In parallel, GPU Interpret Mode got expanded [2] to support while_p operations and better handle unpacking of arbitrary transforms for copy_gmem_to_smem_p, which broadens what workloads can now run in this debugging mode. These are targeted fixes addressing real memory and compatibility pain points from the last cycle.

Action items

→ Sync JAX to latest if you hit triangular_solve OOMs google/jax [immediate]
→ Test GPU Interpret Mode with while_p workloads if debugging GPU kernels google/jax [plan]

References

[1] restore triangular_solve batching rule, improve sharding rule ↗ google/jax
[2] Support while_p in GPU Interpret Mode. ↗ google/jax

Quick answers

What shipped in Google on May 31, 2026?: A partial revert of yesterday's batching rule is rolling back OOMs in triangular_solve while keeping the performance gains. In total, 4 commits and 2 pull requests landed.
Who contributed to Google on May 31, 2026?: 2 developers shipped this update, including mattjj and copybara-service[bot].
What were the notable Google updates?: restore triangular_solve batching rule, improve sharding rule and Support while_p in GPU Interpret Mode.

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.