Do not allow None axis names in meshes.

PiperOrigin-RevId: 686557025

Author: bartchr808

jax/_src/mesh.py

@@ -165,6 +165,8 @@ def __new__(cls, devices: np.ndarray | Sequence[xc.Device],
     if isinstance(axis_names, str):
       axis_names = (axis_names,)
     axis_names = tuple(axis_names)
+    if not all(i is not None for i in axis_names):
+      raise ValueError(f"Mesh axis names cannot be None. Got: {axis_names}")
 
     if devices.ndim != len(axis_names):
       raise ValueError(

tests/array_test.py

@@ -1295,6 +1295,10 @@ def test_to_device(self):
     self.assertEqual(x_device.device, device)
     self.assertEqual(x_sharding.device, sharding)
 
+  def test_mesh_with_axis_name_none(self):
+    with self.assertRaisesRegex(ValueError, 'Mesh axis names cannot be None.'):
+      jax.sharding.Mesh(jax.devices(), (None, 'x'))
+
 
 @jtu.with_config(jax_use_shardy_partitioner=True)
 class ShardyShardingTest(jtu.JaxTestCase):

tests/pallas/tpu_pallas_distributed_test.py

@@ -294,8 +294,8 @@ def test_kernel(x_ref,
             )
         )
 
-    devices = mesh_utils.create_device_mesh((1, num_devices))
-    mesh = jax.sharding.Mesh(devices, P(None, 'x'))
+    devices = mesh_utils.create_device_mesh((num_devices,))
+    mesh = jax.sharding.Mesh(devices, 'x')
     sharding = jax.sharding.NamedSharding(mesh, P(None, 'x'))
     unsharded_arr = jax.random.normal(
         jax.random.key(0), shape=(8, 128 * num_devices))

tests/pallas/tpu_pallas_test.py

@@ -2165,8 +2165,8 @@ def kernel(x_ref, o_ref, send_sem, recv_sem):
       )
     with self.assertRaisesRegex(
         Exception, 'DMAs with bool dtypes are not supported.'):
-      devices = mesh_utils.create_device_mesh((1, num_devices))
-      mesh = jax.sharding.Mesh(devices, P(None, 'x'))
+      devices = mesh_utils.create_device_mesh((num_devices,))
+      mesh = jax.sharding.Mesh(devices, ('x',))
       sharding = jax.sharding.NamedSharding(mesh, P(None, 'x'))
       input_arr = jax.device_put(input_arr, sharding)
       jax.jit(