Support transfer guard in broadcast_one_to_all(). Fixes jax-ml#25325

PiperOrigin-RevId: 703666450

Author: danijar

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(