Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.

This enables CPU collectives by default, making multi-process CPU
communication work without extra configuration.

PiperOrigin-RevId: 722209646

CHANGELOG.md

@@ -25,6 +25,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 * Changes
   * `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` now work as
     env vars. Before they could only be specified via jax.config or flags.
+  * `JAX_CPU_COLLECTIVES_IMPLEMENTATION` now defaults to `'gloo'`, meaning
+    multi-process CPU communication works out-of-the-box.
 
 ## jax 0.5.0 (Jan 17, 2025)
 

jax/_src/xla_bridge.py

@@ -62,6 +62,8 @@
 
 MIN_COMPUTE_CAPABILITY = 52
 
+_DEFAULT_CPU_COLLECTIVES_IMPL = 'gloo'
+
 # TODO(phawkins): Remove jax_xla_backend.
 _XLA_BACKEND = config.string_flag(
     'jax_xla_backend', '',
@@ -244,6 +246,9 @@ def make_cpu_client(
                       '"jax_cpu_collectives_implementation", "gloo")` instead.',
                       DeprecationWarning,
                       )
+    if collectives_impl is None:
+      collectives_impl = _DEFAULT_CPU_COLLECTIVES_IMPL
+
     if collectives_impl == 'gloo':
       collectives = xla_client._xla.make_gloo_tcp_collectives(
         distributed_client=distributed.global_state.client,