Support device_put of an uncommitted array to a single global device in McJAX.

PiperOrigin-RevId: 861417146

Author: emilyfertig

jax/_src/dispatch.py

@@ -518,11 +518,19 @@ def _device_put_sharding_impl(
         return _DeferredShardArg(x, x_sharding, aval, x.committed, copy)
     elif is_single_device_sharding(x_sharding):
       device = x_sharding._device_assignment[0] if device is None else device
+      sharding = SingleDeviceSharding(device)
+      if not x._committed and not sharding.has_addressable_devices:
+        # For uncommitted arrays in McJAX, each process has a local copy of the
+        # array. If the destination sharding is not addressable, no data
+        # transfer is needed, since the data was transferred in the process
+        # in which the sharding is addressable.
+        shards, devices = [], []
+      else:
+        shards, devices = [x], [device]
       if copy == ArrayCopySemantics.ALWAYS_COPY:
-        return xc.batched_device_put(aval, SingleDeviceSharding(device), [x],
-                                     [device], True, True)
-      return pxla.batched_device_put(aval, SingleDeviceSharding(device), [x],
-                                     [device])
+        return xc.batched_device_put(aval, sharding, shards, devices, True,
+                                     True)
+      return pxla.batched_device_put(aval, sharding, shards, devices)
 
   sh = SingleDeviceSharding(pxla.get_default_device()
                             if device is None else device)

tests/multiprocess/array_test.py

@@ -943,9 +943,13 @@ def test_cross_host_transfer_cpu_error(self):
         ValueError, "does not support cross-host device transfers"):
       jax.device_put(y, dst_sharding)
 
+  @parameterized.named_parameters(
+      ("numpy", np.arange),
+      ("uncommitted", jnp.arange),
+  )
   @jtu.skip_on_devices("cpu")
-  def test_cross_host_transfer_single_device_sharding(self):
-    x = np.arange(64).reshape(8, 8)
+  def test_cross_host_transfer_single_device_sharding(self, arange_fn):
+    x = arange_fn(64).reshape(8, 8)
     src_pid = 0
     dst_pid = 1
     src_sharding = jax.sharding.SingleDeviceSharding(
@@ -960,9 +964,13 @@ def test_cross_host_transfer_single_device_sharding(self):
     else:
       self.assertEmpty(z.addressable_shards)
 
+  @parameterized.named_parameters(
+      ("numpy", np.arange),
+      ("uncommitted", jnp.arange),
+  )
   @jtu.skip_on_devices("cpu")
-  def test_cross_host_transfer_named_sharding(self):
-    x = np.arange(64).reshape(8, 8)
+  def test_cross_host_transfer_named_sharding(self, arange_fn):
+    x = arange_fn(64).reshape(8, 8)
     n_local = jax.local_device_count()
     src_pid = 0
     dst_pid = 1
@@ -1113,11 +1121,15 @@ def test_device_put_with_mixed_local_and_remote_transfers(self):
     for shard in z.addressable_shards:
       np.testing.assert_array_equal(shard.data, x[shard.index])
 
+  @parameterized.named_parameters(
+      ("numpy", np.arange),
+      ("uncommitted", jnp.arange),
+  )
   @jtu.skip_on_devices("cpu")
-  def test_device_put_to_device(self):
+  def test_device_put_to_device(self, arange_fn):
     if jaxlib_extension_version < 400:
       self.skipTest("This functionality is not yet supported in jaxlib.")
-    x = np.arange(64).reshape(8, 8)
+    x = arange_fn(64).reshape(8, 8)
     src_pid = 0
     dst_pid = 1
     src_device = jax.local_devices(process_index=src_pid)[0]