[sharding_in_types] Check out_avals with mesh context too. This is because users can pass their own shardings to functions like einsum, reshape, broadcast`, etc

PiperOrigin-RevId: 707672801

Author: yashk2810

jax/_src/lax/utils.py

@@ -54,11 +54,13 @@ def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
   least_specialized = type(max(avals, key=_get_array_abstraction_level))
   if least_specialized is core.ShapedArray:
     core.check_avals_context_mesh(avals, prim.name)
-    return core.ShapedArray(
+    out_aval = core.ShapedArray(
         shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
         weak_type=weak_type,
         sharding=(sharding_rule(*avals, **kwargs)
                   if config.sharding_in_types.value else None))
+    core.check_avals_context_mesh([out_aval], prim.name)
+    return out_aval
   elif least_specialized is core.DShapedArray:
     shape = shape_rule(*avals, **kwargs)
     ty = (core.ShapedArray if all(type(d) is int for d in shape)
@@ -83,9 +85,11 @@ def standard_multi_result_abstract_eval(
     out_shardings = (sharding_rule(*avals, **kwargs)
                      if config.sharding_in_types.value else
                      [None] * len(out_shapes))
-    return [core.ShapedArray(s, d, weak_type=weak_type, sharding=sh)
-            for s, d, weak_type, sh in zip(out_shapes, out_dtypes, weak_types,
-                                           out_shardings)]
+    out_avals = [core.ShapedArray(s, d, weak_type=weak_type, sharding=sh)
+                 for s, d, weak_type, sh in zip(out_shapes, out_dtypes,
+                                                weak_types, out_shardings)]
+    core.check_avals_context_mesh(out_avals, prim.name)
+    return out_avals
   elif least_specialized is core.UnshapedArray:
     out_dtypes = dtype_rule(*avals, **kwargs)
     return [core.UnshapedArray(dtype, weak_type=weak_type)

tests/pjit_test.py

@@ -5827,6 +5827,24 @@ def f(x, y):
     self.assertDictEqual(out1.sharding.mesh.axis_types, {AxisTypes.User: 'x'})
     self.assertDictEqual(out2.sharding.mesh.axis_types, {AxisTypes.Auto: 'x'})
 
+  @jtu.with_user_mesh((2,), 'x')
+  def test_output_different_context_error(self, mesh):
+    np_inp1 = np.arange(16).reshape(8, 2)
+    arr1 = jax.device_put(np_inp1, NamedSharding(mesh, P('x', None)))
+    arr2 = jax.device_put(np_inp1.T, NamedSharding(mesh, P(None, 'x')))
+    auto_mesh = jax.make_mesh((2,), 'x',
+                              axis_types={AxisTypes.Auto: 'x'}).abstract_mesh
+
+    @jax.jit
+    def f(x, y):
+      out = jnp.einsum('xy,yz->xz', x, y,
+                       out_type=NamedSharding(auto_mesh, P('x', None)))
+      return out
+
+    with self.assertRaisesRegex(
+        ValueError, "context mesh.* should match the aval mesh"):
+      f(arr1, arr2)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):