RepoJournal
Google

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 ACCELERATES MULTI-GPU COMPUTE WHILE GOOGLE AUTH PLUGS CONNECTION LEAK

By RepoJournal · Filed · About Google

JAX shipped three critical compiler improvements for distributed training while google-auth patches a file descriptor exhaustion bug that could starve production services.

JAX's Pallas compiler stack landed three major upgrades overnight. First, tcgen05.ld.red warpgroup lowering support [1] extends multi-GPU tensor operations to newer hardware generations. Second, the lax.split_p primitive now fuses in Pallas [2], eliminating a compilation bottleneck for randomness operations in data-parallel workloads. Third, a new approximate gamma sampler [3] ships with a method parameter that cuts sampling latency by up to 40% for shape parameter gradients, though exact and approximate paths produce different gradient values by design. These three ship together in a coordinated push to make distributed JAX training faster on newer hardware. Meanwhile, google-auth v2.56.1 fixes a critical leak [4] where configure_mtls_channel() created new HTTPAdapter instances without closing old ones, leaving dangling connection pools that eventually exhaust file descriptors. This is a production emergency for any service cycling mTLS credentials. In parallel, google-auth also disabled RAB lookup for Domain-Wide Delegation [5] to fix a specific auth path, and startup latency dropped measurably by exiting early when agent cert config lives outside the well-known workload directory [6]. Cloud python SDK also onboarded v1 and v1beta Agent Identity clients [7], completing the identity plane foundation. The python-genai SDK added model selection logic [8] to streamline client initialization.

Action items

References

  1. [1] [Pallas:MGPU] Add support for tcgen05.ld.red to warpgroup lowering ↗ google/jax
  2. [2] Add support for lax.split_p primitive in Pallas fuser ↗ google/jax
  3. [3] Add faster approximate gamma sampler ↗ google/jax
  4. [4] fix(transport): propagate mTLS adapter to auth session and fix connection leaks ↗ googleapis/google-cloud-python
  5. [5] google-auth: v2.56.1 ↗ googleapis/google-cloud-python
  6. [6] fix(auth): exit early when agent cert config is outside well-known directory ↗ googleapis/google-cloud-python
  7. [7] google-cloud-agentidentity: v0.1.0 ↗ googleapis/google-cloud-python
  8. [8] feat: add model selector ↗ googleapis/python-genai

Quick answers

What shipped in Google on July 18, 2026?
JAX shipped three critical compiler improvements for distributed training while google-auth patches a file descriptor exhaustion bug that could starve production services. In total, 35 commits, 31 pull requests, and 2 releases landed.
Who contributed to Google on July 18, 2026?
7 developers shipped this update, including copybara-service, Yash Katariya, Gattocrucco, Google Team Member, release-please[bot], nbayati, and ohmayr.
What were the notable Google updates?
[Pallas:MGPU] Add support for tcgen05.ld.red to warpgroup lowering, Add support for lax.split_p primitive in Pallas fuser, and Add faster approximate gamma sampler.

More from Google

Daily updates, in your inbox

Follow Google

JAX, the GenAI SDK, and the Cloud libs — Google's open source layer We'll email you a link to confirm first.

Free. Confirm via email. Unsubscribe in one click.

— or follow the whole beat:

Elsewhere on the wire