$ the-wire · showcase
JAX removes deprecated Pallas APIs, adds barrier-skip for multi-device calls
By RepoJournal · Filed · About Google
JAX purged a set of deprecated Pallas APIs and gave Mosaic GPU callers a way to skip multi-device barriers.
JAX removed several previously deprecated Pallas APIs as part of the regular deprecation cycle. Code that still calls these removed APIs will now raise an AttributeError; migrate to the current replacements before upgrading [1].
Mosaic GPU's pl.pallas_call now accepts a skip_device_barrier metadata parameter. Setting it causes the custom_call to skip LaunchMultiGpuBarrier, which can reduce synchronization overhead for single-device or already-synchronized workloads [2].
Python 3.15 support work continues: TPU CI no longer checks libtpu.so strings because the hardcoded path does not exist on Python 3.15, and the obsolete cloud-tpu-ci-nightly.yml workflow was deleted [3][4]. A concurrent change guards PyArray_Storage::dynamic_shape with the storage mutex and encapsulates the computation in PyArray::dynamic_shape() [5].
The google-cloud-python ecosystem adds Cloud Spanner Queues send and ack mutations [6], while JAX's auth layer now parses hostnames with urllib.parse.urlsplit to avoid false positives when deciding mTLS, and adds support for PSC custom domains [7].
Action items
- → If you call deprecated Pallas APIs, migrate to the current replacements before upgrading JAX google/jax [plan]
- → For Mosaic GPU workloads that don't need multi-device barriers, set skip_device_barrier metadata to cut sync overhead google/jax [plan]
- → If you rely on Cloud Spanner Queues, test the new send and ack mutations googleapis/google-cloud-python [monitor]
References
- [1] [pallas] Removed a few previously deprecated APIs ↗ google/jax
- [2] [Mosaic GPU] Add parameter to skip multi-device barriers passed through metadata. ↗ google/jax
- [3] Remove libtpu.so strings check from TPU CI scripts. ↗ google/jax
- [4] Remove libtpu.so strings check from TPU CI scripts. ↗ google/jax
- [5] Guard PyArray_Storage::dynamic_shape with mu ↗ google/jax
- [6] feat: add send and ack mutations Cloud Spanner Queues ↗ googleapis/google-cloud-python
- [7] fix(auth): parse hostname for mTLS and PSC endpoint certificate rotation (#18147) ↗ googleapis/google-cloud-python