Fix device_put's abstract_eval to return the correct sharding on the type when the sharding passed to it changes. Fixes #31793

PiperOrigin-RevId: 807429077

Author: yashk2810

jax/_src/api.py

@@ -2719,13 +2719,16 @@ def device_put(
         assert not m and not d
         copy_semantics.append(dispatch.ArrayCopySemantics.ALWAYS_COPY)
 
-    x_avals = tuple(shaped_abstractify(i) for i in x_flat)
-    for aval, d in zip(x_avals, device_flat):
+    dst_avals = []
+    for xf, d in zip(x_flat, device_flat):
+      aval = shaped_abstractify(xf)
+      aval = dispatch.update_dp_aval(aval, d)
+      dst_avals.append(aval)
       _check_sharding(aval, d)
     if core.trace_state_clean():
       out_flat = dispatch._batched_device_put_impl(
           *x_flat, devices=device_flat, srcs=src_flat,  # type: ignore
-          copy_semantics=copy_semantics, x_avals=x_avals)
+          copy_semantics=copy_semantics, dst_avals=dst_avals)
     else:
       out_flat = dispatch.device_put_p.bind(
           *x_flat, devices=tuple(device_flat), srcs=tuple(src_flat),

jax/_src/dispatch.py

@@ -519,6 +519,7 @@ def _device_put_impl(
   if aval is None:
     try:
       aval = core.abstractify(x)
+      aval = update_dp_aval(aval, device)
     except TypeError as err:
       raise TypeError(
           f"Argument '{x}' of type {type(x)} is not a valid JAX type") from err
@@ -557,11 +558,11 @@ def _batched_device_put_impl(
     devices: Sequence[Device | Sharding | Format | None],
     srcs: Sequence[Device | Sharding | Format | None],
     copy_semantics: Sequence[ArrayCopySemantics],
-    x_avals: Sequence[core.ShapedArray | None]):
+    dst_avals: Sequence[core.ShapedArray | None]):
   ys = []
   dsa_indices, dsa_xs, dsa_shardings, dsa_copy_semantics = [], [], [], []
   for i, (x, device, src, cp, aval) in enumerate(
-      zip(xs, devices, srcs, copy_semantics, x_avals)):
+      zip(xs, devices, srcs, copy_semantics, dst_avals)):
     y = _device_put_impl(x, device=device, src=src, copy=cp, aval=aval)
     if isinstance(y, _DeferredShardArg):
       dsa_indices.append(i)
@@ -589,24 +590,29 @@ def batched_device_put_impl(
     copy_semantics: Sequence[ArrayCopySemantics]):
   return _batched_device_put_impl(
       *xs, devices=devices, srcs=srcs, copy_semantics=copy_semantics,
-      x_avals=[None] * len(devices))
+      dst_avals=[None] * len(devices))
 
 
 device_put_p = core.Primitive('device_put')
 device_put_p.multiple_results = True
 device_put_p.def_impl(batched_device_put_impl)
 
 
+def update_dp_aval(aval, d):
+  if not isinstance(aval, core.ShapedArray):
+    return aval
+  if isinstance(d, Sharding):
+    aval = (aval.update(sharding=aval.sharding.update(mesh=d.mesh.abstract_mesh))
+            if isinstance(d, NamedSharding) else aval.update(sharding=None))
+    if d.memory_kind is not None:
+      aval = aval.update(memory_space=core.mem_kind_to_space(d.memory_kind))
+    return aval
+  elif isinstance(d, core.MemorySpace):
+    return aval.update(memory_space=d)
+  return aval
+
 def _device_put_abstract_eval(*xs, devices, srcs, copy_semantics):
-  out = []
-  for x, d in zip(xs, devices):
-    if isinstance(d, Sharding) and d.memory_kind is not None:
-      out.append(x.update(memory_space=core.mem_kind_to_space(d.memory_kind)))
-    elif isinstance(d, core.MemorySpace):
-      out.append(x.update(memory_space=d))
-    else:
-      out.append(x)
-  return out
+  return [update_dp_aval(x, d) for x, d in zip(xs, devices)]
 device_put_p.def_abstract_eval(_device_put_abstract_eval)
 
 def _device_put_transpose(cts, *_, devices, srcs, copy_semantics):

tests/pjit_test.py

@@ -8894,6 +8894,15 @@ def model(kernel, inputs):
     self.assertEqual(out1.sharding, kernel.sharding)
     self.assertEqual(out2.sharding, inputs.sharding)
 
+  @jtu.with_explicit_mesh((2,), 'x')
+  def test_device_put_typeof(self, mesh):
+    array = jnp.zeros(8)
+    self.assertEqual(jax.typeof(array).sharding,
+                     NamedSharding(mesh.abstract_mesh, P(None)))
+
+    array = jax.device_put(array, SingleDeviceSharding(jax.devices()[0]))
+    self.assertTrue(jax.typeof(array).sharding.mesh.empty)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):