RepoJournal
Google

Google

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

Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.

One email a day. Unsubscribe in one click.

Pick a date

Topics: AI / ML Python Full archive →

The Wire · Showcase

JAX SWITCHES TO BZLMOD, SHIPS AOT COMPILATION FOR TRITON KERNELS

By RepoJournal · Filed · About Google

JAX's build system migration to Bzlmod is live, paired with major performance improvements for autotuned GPU kernels that compile at instantiation rather than runtime.

The JAX team switched to Bzlmod by default [1], a significant infrastructure move that restructures how pip dependencies are exposed in the build system. Pip packages now route through `@pypi//<package>` rather than standalone repositories, requiring updates to lockfiles and dependency declarations. Simultaneously, JAX shipped AOT compilation for autotuned Triton kernels [2], moving compilation of candidate configurations from JIT time into the instantiate phase. This eliminates serialization overhead since instantiate and execute now happen in the same process. Supporting infrastructure landed across three commits: Triton now handles host-side TMA descriptors [4], the pallas multi-GPU layer auto-infers `PROMISE_IN_BOUNDS` in pipeline emission [5], and requirements lock files were regenerated [3] to accommodate the Bzlmod transition. On the GenAI front, the Python SDK cleaned deprecated Turn types [7], dropped unused auth imports [6], and fixed automatic function calling bugs in streaming contexts [8]. BigFrames hit v2.48.0 with expanded string operation support in the transpiler [9], while BigTable standardized client-side metrics across implementations [10]. The google-cloud-python repo saw two reverts [11] [12] of generator and librarian changes, suggesting recent commits needed correction.

One email a day. Unsubscribe in one click.

Action items

References

  1. [1] Switch to bzlmod by default. ↗ google/jax
  2. [2] Implement AOT compilation for autotuned JAX-Triton kernels. ↗ google/jax
  3. [3] Regenerate requirements lock files. ↗ google/jax
  4. [4] [triton] Added support for host-side TMA descriptors ↗ google/jax
  5. [5] [pallas:mgpu] Automatically infer `PROMISE_IN_BOUNDS` in `emit_pipeline` ↗ google/jax
  6. [6] chore: drop unused auth imports from the MCP utils module googleapis/python-genai
  7. [7] chore: remove deprecated Turn types googleapis/python-genai
  8. [8] fix: improve AFC(automatic function calling) in chat including fixing bugs in AFC for generate_content_stream. also log warnings in generate_content, generate_content_stream and their async variants t ↗ googleapis/python-genai
  9. [9] bigframes: v2.48.0 ↗ googleapis/google-cloud-python
  10. [10] fix(bigtable): standardize client side metrics (#17899) googleapis/google-cloud-python
  11. [11] Revert "fix(generator): use flat_ref_types in test templates and delete remove-unused-imports" (#18088) googleapis/google-cloud-python
  12. [12] Revert "chore(librarian): remove post processing script for ces" (#18094) googleapis/google-cloud-python

Quick answers

What shipped in Google on August 13, 2026?
JAX's build system migration to Bzlmod is live, paired with major performance improvements for autotuned GPU kernels that compile at instantiation rather than runtime. In total, 49 commits, 44 pull requests, and 2 releases landed.
Who contributed to Google on August 13, 2026?
7 developers shipped this update, including hawkinsp, copybara-service, Google Team Member, Annie Luc, release-please[bot], Daniel Sanche, and ohmayr.
What were the notable Google updates?
Switch to bzlmod by default, Implement AOT compilation for autotuned JAX-Triton kernels, and Regenerate requirements lock files.

More from Google

Daily updates, in your inbox

Follow Google

Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.

or

One email a day. Unsubscribe in one click. Read a past issue →

Elsewhere on the wire

Want every project, not just this one?