Allow eval_shape to propagate shardings if the aval has shardings in full explicit mode

PiperOrigin-RevId: 761708753

Author: yashk2810

jax/_src/pjit.py

@@ -350,11 +350,17 @@ def jit_lower(jit_func, *args, **kwargs):
 @api_boundary
 def jit_eval_shape(jit_func, *args, **kwargs):
   p, _ = _infer_params(jit_func._fun, jit_func._jit_info, args, kwargs)
-  out_s = [None if isinstance(s, UnspecifiedValue) else s for s in p.params['out_shardings']]
-  # TODO(yashkatariya): Add `Layout` to SDS.
-  out = [api.ShapeDtypeStruct(x.shape, x.dtype, sharding=s,
-                              weak_type=x.weak_type)
-         for x, s in zip(p.params['jaxpr'].out_avals, out_s)]
+  out_shardings = [None if isinstance(s, UnspecifiedValue) else s
+                   for s in p.params['out_shardings']]
+  out = []
+  for a, out_s in zip(p.params['jaxpr'].out_avals, out_shardings):
+    if out_s is None:
+      s = a.sharding if a.sharding.mesh._are_all_axes_explicit else out_s
+    else:
+      s = out_s
+    # TODO(yashkatariya): Add `Layout` to SDS.
+    out.append(api.ShapeDtypeStruct(a.shape, a.dtype, sharding=s,
+                                    weak_type=a.weak_type))
   return tree_unflatten(p.out_tree, out)
 
 def jit_evict_fn(self):

tests/pjit_test.py

@@ -7857,6 +7857,20 @@ def g(x, y):
         core.ShardingTypeError, "lhs is unreduced while rhs is not"):
       g.trace(x, y)
 
+  @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
+  def test_eval_shape(self, mesh):
+    np_inp = np.arange(16).reshape(8, 2)
+    arr = jax.device_put(np_inp, P('x', 'y'))
+
+    @jax.jit
+    def f(x):
+      return x * 2
+
+    out = jax.eval_shape(f, arr)
+    self.assertIsInstance(out, jax.ShapeDtypeStruct)
+    self.assertEqual(out.sharding,
+                     NamedSharding(mesh.abstract_mesh, P('x', 'y')))
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):