The Wire · Showcase
JAX hardens distributed mode with mTLS
By RepoJournal · Filed · About Google
JAX just made its distributed coordination service production-ready for untrusted networks.
JAX's distributed coordination service now supports mutual TLS, closing a gap where OSS `jax.distributed.initialize` only handled insecure credentials [1]. The change also plumbs `verify_secure_credentials` through the API and adds a Security considerations docs page. In the same push, the codebase is renaming `VJPHiPrimitive` to `HiPrim` with a backward-compatible alias [2], then migrating all usages and deleting the old name [3]. That's a breaking change if you were deep in the internal higher-order primitive ecosystem, but the alias softens the landing. Meanwhile, mattjj fixed a class of tracing bugs across higher-order APIs by switching to the tracing pytree registry [4]. google-cloud-python is quieter: contributor docs are being centralized so package-level `CONTRIBUTING.rst` files just point to the repo root [5][6], and the docfx CI job moves to post-submit [7].
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
- → Review your JAX distributed setup and adopt mTLS credentials where available google/jax [plan]
- → Check for usage of VJPHiPrimitive and prepare for the eventual HiPrim migration google/jax [monitor]
- → Update your contributing docs to point to the centralized root CONTRIBUTING.rst googleapis/google-cloud-python [plan]
References
- [1] Add mTLS support to `jax.distributed` coordination service google/jax
- [2] Rename `VJPHiPrimitive` to `HiPrim` but keep `VJPHiPrimitive` as an alias ↗ google/jax
- [3] Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive ↗ google/jax
- [4] use tracing registry in other higher-order APIs, fix explain bugs ↗ google/jax
- [5] docs(handwritten): centralize CONTRIBUTING.rst pointers ↗ googleapis/google-cloud-python
- [6] docs(handwritten): centralize CONTRIBUTING.rst pointers (#17642) googleapis/google-cloud-python
- [7] ci: move docfx job to post-submit job (#18182) googleapis/google-cloud-python