Error in jax.make_mesh if we detect multi-slice topology

PiperOrigin-RevId: 829634731

Author: yashk2810

jax/_src/sharding_impls.py

@@ -1191,6 +1191,11 @@ def make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str],
   mesh_devices = mesh_utils.create_device_mesh(
       new_axis_shapes, devices,
       allow_split_physical_axes=allow_split_physical_axes)
+  if (hasattr(mesh_devices.flat[0], 'slice_index') and
+      len({d.slice_index for d in mesh_devices.flat}) > 1):
+    raise ValueError(
+        '`jax.make_mesh` does not support multi-slice topologies. Please use'
+        ' jax.experimental.mesh_utils.create_hybrid_device_mesh')
   if axis_types is None:
     if deprecations.is_accelerated('jax-make-mesh-default-explicit'):
       axis_types = (mesh_lib.AxisType.Explicit,) * len(mesh_devices.shape)