Support transfer guard in broadcast_one_to_all(). Fixes #25325

PiperOrigin-RevId: 703666450
jax-ml · Dec 7, 2024 · 861115a · 861115a
1 parent baedb62
commit 861115a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/experimental/multihost_utils.py b/jax/experimental/multihost_utils.py
@@ -75,7 +75,7 @@ def pre_jit(x):
     return host_local_array_to_global_array(inp, global_mesh, pspec)
 
   def post_jit(x):
-    return np.asarray(x.addressable_data(0))
+    return jax.device_get(x.addressable_data(0))
 
   in_tree = jax.tree.map(pre_jit, in_tree)
   out_tree = jax.jit(_psum, out_shardings=jax.sharding.NamedSharding(