Repository: jax-ml/jax
Pull request #27485 (2025-03-26 18:47:27)

Raise an informative error when the length of device_assignment doesn't match the mesh.size of out_avals. This happens when (1) we can't extract the device_assignment from the arguments and (2) there is no concrete mesh in context.
Raise an informative error when the length of device_assignment doesn't match the mesh.size of out_avals. This happens when (1) we can't extract the device_assignment from the arguments and (2) there is no concrete mesh in context.

For example:

```
def test_random_normal_wo_mesh_context_error(self):
    mesh = jtu.create_mesh((2, 2), ('x', 'y'),
                           axis_types=(AxisType.Explicit,) * 2)
    s = NamedSharding(mesh, P('x', 'y'))

    @jax.jit
    def f(key):
      out = jax.random.normal(key, shape=(8, 12), out_sharding=s)
      self.assertEqual(out.aval.sharding.spec, P('x', 'y'))
      self.assertEqual(out.aval.sharding.mesh, mesh.abstract_mesh)
      return out

    key = jax.random.key(1)
    with self.assertRaisesRegex(
        ValueError,
        'Length of device assignment.*is not equal to the size of the mesh'):
      f(key)
```



Patch:
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index 387f0661a..6f95b1b72 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2281,6 +2281,16 @@ def lower_sharding_computation(
       devices_from_context)
   unique_intermediate_shardings = [js for js, _ in unique_intermediate_shardings]
 
+  for a in global_out_avals:
+    if (a is not core.abstract_token and not a.sharding.mesh.empty and
+        a.sharding.mesh._are_all_axes_explicit and
+        len(device_assignment) != a.sharding.mesh.size):
+      raise ValueError(
+          f"Length of device assignment {len(device_assignment)} is not equal"
+          f" to the size of the mesh {a.sharding.mesh.size} of aval"
+          f" {a.str_short(True, True)}. Please enter your `jit` into a mesh"
+          " context via `jax.sharding.use_mesh`.")
+
   # TODO(parkers): One _raw_platform has been unified with platform,
   # change this back to just read platform.
   platforms = lowering_platforms or (


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 5cd16e1e6..277f24bd7 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7265,6 +7265,24 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     else:
       self.assertIn('mhlo.sharding = "{devices=[2,2]<=[4]}"}', lowered_text)
 
+  def test_random_normal_wo_mesh_context_error(self):
+    mesh = jtu.create_mesh((2, 2), ('x', 'y'),
+                           axis_types=(AxisType.Explicit,) * 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+
+    @jax.jit
+    def f(key):
+      out = jax.random.normal(key, shape=(8, 12), out_sharding=s)
+      self.assertEqual(out.aval.sharding.spec, P('x', 'y'))
+      self.assertEqual(out.aval.sharding.mesh, mesh.abstract_mesh)
+      return out
+
+    key = jax.random.key(1)
+    with self.assertRaisesRegex(
+        ValueError,
+        'Length of device assignment.*is not equal to the size of the mesh'):
+      f(key)
+
   def test_random_normal_wo_mesh_context(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'),
                            axis_types=(AxisType.Explicit,) * 2)