RepoJournal
Google

Google

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

Pick a date

  1. Sun 3 may
  2. Mon 4 may
  3. Tue 5 may
  4. Wed 6 may
  5. Thu 7 may
  6. Fri 8 may
  7. Sat 9 may
  8. Sun 10 may
  9. Mon 11 may
  10. Tue 12 may
  11. Wed 13 may
  12. Fri 15 may
  13. Sat 16 may
  14. Sun 17 may
  15. Mon 18 may
  16. Tue 19 may
  17. Wed 20 may
  18. Thu 21 may
  19. Sat 23 may
  20. Sun 24 may
  21. Mon 25 may
  22. Tue 26 may
  23. Thu 28 may
  24. Fri 29 may
  25. Sat 30 may
  26. Sun 31 may
  27. Mon 1 jun
  28. Mon 8 jun
  29. Tue 9 jun
  30. Today 10 jun

The Wire · Showcase

JAX FIXES CACHE KEY BUGS BREAKING DISTRIBUTED TRAINING ACROSS ENVIRONMENTS

By RepoJournal · Filed · About Google

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

References

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

FAQ

What changed 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.
What should Google teams do about it?
Pin JAX to the latest commit with cache key fix before scaling distributed jobs • Update google-cloud-python past version supporting Python 3.7-3.9 removal • If using Pallas multi-GPU kernels, migrate pl.program_id to lax.axis_index
Which Google repositories shipped on May 23, 2026?
google/jax, googleapis/google-cloud-python, googleapis/python-genai

Related across the cluster

For your repos

The showcase is a teaser.
Your wire is the product.

Same engine. Different stack. Below: what changes when the wire is yours.

Showcase wire

  • •14 famous open source orgs
  • •One wire per day
  • •Public, generic
  • •Read on the web, when you remember

Your wire

  • →Up to 1,500 of your repos - orgs, deps, vendors
  • →Morning and evening briefs
  • →Action items routed to your team
  • →Slack delivery, email, breaking-news CVE alerts

Want a hands-on demo first? Ask a current user for an invite link.