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

Reverts a158e02

PiperOrigin-RevId: 701045239

Author: yashk2810

jax/_src/mesh_utils.py

@@ -705,6 +705,15 @@ def _transpose_trick(
       *_TRANSPOSE_TRICKS[topology][mesh_shape_no_trivial_dims]
   )
 
+def _canonicalize_axis_sizes(axis_sizes: Sequence[int]
+                             ) -> tuple[int, ...] | None:
+  new_sizes = []
+  for s in axis_sizes:
+    try:
+      new_sizes.append(int(s))
+    except:
+      return None
+  return tuple(new_sizes)
 
 def create_device_mesh(
     mesh_shape: Sequence[int],
@@ -740,33 +749,41 @@ def create_device_mesh(
   """
   if devices is None:
     devices = xb.devices()
-  if np.prod(mesh_shape) != len(devices):
+
+  new_mesh_shape = _canonicalize_axis_sizes(mesh_shape)
+  if new_mesh_shape is None:
+    raise ValueError(
+        f'`mesh_shape` passed to `create_device_mesh` should be a sequence of'
+        f' ints. Got {mesh_shape}')
+  del mesh_shape
+
+  if math.prod(new_mesh_shape) != len(devices):
     raise ValueError(
         f'Number of devices {len(devices)} must equal the product '
-        f'of mesh_shape {mesh_shape}'
+        f'of mesh_shape {new_mesh_shape}'
     )
   last_device = devices[-1]
 
   handler = device_kind_handler_dict.get(last_device.device_kind, None)
   if handler is not None:
     result = handler(
-        mesh_shape, devices, contiguous_submeshes=contiguous_submeshes
+        new_mesh_shape, devices, contiguous_submeshes=contiguous_submeshes
     )
     if result is not None:
       return result
 
   if last_device.platform == 'tpu':
     physical_mesh = _get_physical_tpu_mesh(devices)
     if contiguous_submeshes:
-      physical_mesh = _transpose_trick(physical_mesh, mesh_shape)
+      physical_mesh = _transpose_trick(physical_mesh, new_mesh_shape)
     device_mesh, _ = _create_device_mesh_for_nd_torus(
         physical_mesh,
-        mesh_shape,
+        new_mesh_shape,
         allow_split_physical_axes=allow_split_physical_axes,
     )
     return device_mesh
   else:
-    device_mesh = np.asarray(devices).reshape(mesh_shape)
+    device_mesh = np.asarray(devices).reshape(new_mesh_shape)
     return device_mesh
 
 

jax/_src/sharding_impls.py

@@ -1714,17 +1714,25 @@ def make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str],
   """
   if devices is None:
     devices = xla_bridge.devices()
-  axis_size = math.prod(axis_shapes)
+  new_axis_shapes = mesh_utils._canonicalize_axis_sizes(axis_shapes)
+  if new_axis_shapes is None:
+    raise ValueError(
+        '`axis_shapes` passed to `make_mesh` should be a sequence of ints.'
+        f' Got {axis_shapes}')
+  del axis_shapes
+
+  axis_size = math.prod(new_axis_shapes)
   if axis_size > len(devices):
     raise ValueError(
         f'Number of devices {len(devices)} must be >= the product '
-        f'of mesh_shape {axis_shapes}')
+        f'of mesh_shape {new_axis_shapes}')
   elif axis_size < len(devices):
     devices = devices[:axis_size]
   if devices[0].device_kind in (mesh_utils._TPU_V5_LITE, mesh_utils._TPU_V5E):
     allow_split_physical_axes = True
   else:
     allow_split_physical_axes = False
   mesh_devices = mesh_utils.create_device_mesh(
-      axis_shapes, devices, allow_split_physical_axes=allow_split_physical_axes)
+      new_axis_shapes, devices,
+      allow_split_physical_axes=allow_split_physical_axes)
   return mesh_lib.Mesh(mesh_devices, axis_names)

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,14 @@ 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'))
+
+    jax.make_mesh((1, np.int32(1), np.int64(1)), ('x', 'y', 'z'))  # doesn't crash
+
   def test_jnp_array_reshard_error(self):
     if jax.device_count() < 2:
       self.skipTest('Requires >=2 devices')