Who contributed to Google on May 25, 2026?

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

What were the notable Google updates?

Reverts ac12cbaa1f877ef38a128000b8c561f136fc7f01, Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose, and Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose.

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 CRITICAL SHARD_MAP GRADIENT BUG AFTER PARTIAL ROLLBACK

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

1 person shipped this

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

JAX reverted a commit that introduced incorrect gradient computation in distributed operations, then shipped the proper fix with safeguards.

The team rolled back ac12cbaa1f [1] after discovering it was inserting unnecessary collective operations during backward passes in shard_map. The root cause: wrong pspec calculations when check_vma=False [2]. Engineers have now landed the corrected version [3] that properly uses the cotangent spec to determine which mesh axes need synchronization. New test coverage [4] ensures unreduced forward inputs trigger correct psum insertion on the backward pass. This fixes a class of silent correctness bugs in distributed training where gradients would compute but with unnecessary communication overhead.

Action items

→ Pull latest JAX and run distributed training tests to verify gradient correctness google/jax [plan]
→ Check if your shard_map usage with check_vma=False was affected by the reverted commit google/jax [monitor]

References

[1] Reverts ac12cbaa1f877ef38a128000b8c561f136fc7f01 ↗ google/jax
[2] Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose. google/jax
[3] Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose. ↗ google/jax
[4] Add a test to make sure that if input is unreduced on fwd (with check_vma=False on shard_map), shard_map inserts a psum on bwd similar to invarying. ↗ google/jax

Quick answers

What shipped in Google on May 25, 2026?: JAX reverted a commit that introduced incorrect gradient computation in distributed operations, then shipped the proper fix with safeguards. In total, 4 commits and 3 pull requests landed.
Who contributed to Google on May 25, 2026?: 1 developer shipped this update, including copybara-service[bot].
What were the notable Google updates?: Reverts ac12cbaa1f877ef38a128000b8c561f136fc7f01, Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose, and Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose.

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.