Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh

PiperOrigin-RevId: 700779838

Author: yashk2810

jax/_src/mesh_utils.py

@@ -705,6 +705,12 @@ def _transpose_trick(
       *_TRANSPOSE_TRICKS[topology][mesh_shape_no_trivial_dims]
   )
 
+def _validate_axis_shapes(axis_shapes: Sequence[int], arg_name: str,
+                          fun_name: str):
+  if not all(isinstance(s, int) for s in axis_shapes):
+    raise ValueError(
+        f'{arg_name} passed to {fun_name} should be a sequence of ints. Got'
+        f' {axis_shapes}')
 
 def create_device_mesh(
     mesh_shape: Sequence[int],
@@ -740,7 +746,8 @@ def create_device_mesh(
   """
   if devices is None:
     devices = xb.devices()
-  if np.prod(mesh_shape) != len(devices):
+  _validate_axis_shapes(mesh_shape, 'mesh_shape', 'create_device_mesh')
+  if math.prod(mesh_shape) != len(devices):
     raise ValueError(
         f'Number of devices {len(devices)} must equal the product '
         f'of mesh_shape {mesh_shape}'

jax/_src/sharding_impls.py

@@ -1714,6 +1714,7 @@ def make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str],
   """
   if devices is None:
     devices = xla_bridge.devices()
+  mesh_utils._validate_axis_shapes(axis_shapes, 'axis_shapes', 'make_mesh')
   axis_size = math.prod(axis_shapes)
   if axis_size > len(devices):
     raise ValueError(

tests/mesh_utils_test.py

@@ -353,6 +353,12 @@ def test_create_device_mesh_for_nd_torus(
         )
     self.assertArraysEqual(assignment, expected_assignment_matrix)
 
+  def test_create_device_mesh_non_int_error(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        "mesh_shape passed to create_device_mesh should be a sequence of ints"):
+      mesh_utils.create_device_mesh(((4,), 4))
+
   @parameterized.named_parameters(
       ('2x2x1', mock_2x2x1_devices,),
       ('2x2x4', mock_2x2x4_devices, ),

tests/pjit_test.py

@@ -4458,6 +4458,12 @@ def g(x):
     self.assertEqual(out2.sharding, s)
     self.assertEqual(out2.dtype, np.float32)
 
+  def test_make_mesh_non_int_error(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        "axis_shapes passed to make_mesh should be a sequence of ints"):
+      jax.make_mesh(((4,), 4), ('x', 'y'))
+
   def test_jnp_array_reshard_error(self):
     if jax.device_count() < 2:
       self.skipTest('Requires >=2 devices')