From 0560cc478ec6aed7c4549c9ca7e10bdd5d126578 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Wed, 15 Nov 2023 13:26:53 -0800 Subject: [PATCH] [JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu"). An upcoming change to JAX will include non-local (addressable) CPU devices in jax.devices() when JAX is used multicontroller-style, where there are multiple Python processes. This change preserves the current behavior by replacing uses of jax.devices("cpu"), which previously only returned local devices, with jax.local_devices("cpu"), which will return local devices both now and in the future. This change is always be safe (i.e., it should always preserve the previous behavior) but it may sometimes be unnecessary if code is never used in a multicontroller setting. PiperOrigin-RevId: 582786346 --- jax/_src/lax/control_flow/loops.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py index 160bae300..8c581bf27 100644 --- a/jax/_src/lax/control_flow/loops.py +++ b/jax/_src/lax/control_flow/loops.py @@ -697,8 +697,11 @@ def _scan_partial_eval(trace, *tracers, reverse, length, num_consts, num_carry, def _maybe_put(x): if isinstance(x, np.ndarray): return dispatch._put_x( - x, jax.sharding.SingleDeviceSharding(jax.devices('cpu')[0]), - shaped_abstractify(x), False) + x, + jax.sharding.SingleDeviceSharding(jax.local_devices(backend='cpu')[0]), + shaped_abstractify(x), + False, + ) else: return x