The Wire · Showcase
JAX CLEARS THE DECK: NVSHMEM DEPRECATED, NUMPY 2.0 SUPPORT DROPPED
By RepoJournal · Filed · About Google
JAX is shipping breaking changes across the stack today, starting with mandatory migration away from NVSHMEM to NCCL Device API and dropping support for NumPy 2.0 and SciPy 1.14 in the 0.11.0 release.
The Mosaic GPU team is removing NVSHMEM usage from XLA custom kernels, forcing customers onto NCCL Device API through collective metadata [1]. This is part of a broader cleanup push: JAX is dropping the deprecated public API module jax.cloud_tpu_init [4], removing stale JAX_IFRT_VERSION ifdefs now that minimum version sits at 55 [3], and purging version tests since jaxlib's minimum is now v0.10.2 [2]. The 0.11.0 release is the hard cutoff for NumPy 2.0 and SciPy 1.14 support [5]. Over on the Python Genai side, version 2.9.0-rc1 is in flight [6]. Google Cloud Python is tightening its typing game: Firestore's CollectionReference now properly annotates .document() calls to return DocumentReference types instead of inheriting base collection typing [7], while the full library suite just finished a librarian regeneration cycle to sync config with generated code [8].
One email a day. Unsubscribe in one click.
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. Read a past issue →
Action items
- → If you're on Mosaic GPU, audit custom kernels for NVSHMEM usage and migrate to NCCL Device API before JAX 0.11.0 google/jax [immediate]
- → Drop NumPy 2.0 and SciPy 1.14 from your dependency matrices before JAX 0.11.0 release google/jax [plan]
- → Update any code calling jax.cloud_tpu_init - it's being removed google/jax [plan]
References
- [1] [Mosaic:GPU] Remove NVSHMEM usage from Mosaic XLA custom kernel and tests. google/jax
- [2] Remove stale version tests now the minimum jaxlib is v0.10.2. google/jax
- [3] Remove stale JAX_IFRT_VERSION ifdefs. google/jax
- [4] Drop the deprecated public API module jax.cloud_tpu_init. google/jax
- [5] Drop support for NumPy 2.0 and SciPy 1.14 for the jax 0.11.0 release. google/jax
- [6] Copybara import of the project: googleapis/python-genai
- [7] chore(firestore): improve typing for collection (#16999) googleapis/google-cloud-python
- [8] chore: re-generate all (#17505) googleapis/google-cloud-python