Fix memory space propagation with call_exported_p. Fixes #33471

PiperOrigin-RevId: 835253850

Author: yashk2810

jax/_src/export/_export.py

@@ -1405,7 +1405,8 @@ def pp_arg_dim(dim_idx: int | None) -> str:
   out_avals = tuple(
       core.ShapedArray(core.evaluate_shape(out_aval.shape, exported_dim_vars,
                                            *exported_dim_values),
-                       dtype=out_aval.dtype, weak_type=out_aval.weak_type)
+                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,
+                       memory_space=out_aval.memory_space)
       for out_aval in exported.out_avals)
   return out_avals, set(exported.ordered_effects + exported.unordered_effects)
 

tests/memories_test.py

@@ -743,6 +743,23 @@ def test_host_to_device_transfer(self):
       self.assertEqual(d.sharding.memory_kind, 'device')
       self.assertArraysEqual(d, orig)
 
+  def test_memory_space_propagated_identity_jit(self):
+    shd = jax.sharding.SingleDeviceSharding(
+        jax.devices()[0], memory_kind='pinned_host')
+    a = jax.device_put(1, shd)
+
+    f = jax.jit(lambda x: x, out_shardings=shd)
+    b = f(a)
+    self.assertEqual(b.sharding, a.sharding)
+
+    f = jax.jit(lambda x: x)
+    b = f(a)
+    self.assertEqual(b.sharding, a.sharding)
+
+    exported = jax.export.export(f)(a)
+    b = exported.call(a)
+    self.assertEqual(b.sharding, a.sharding)
+
 
 class ComputeOffload(jtu.BufferDonationTestCase):