Repository: google/jax
Pull request #20213 (2024-03-12 23:17:36)

Enable Creating Device Mesh with Physical Axes Splits.
Enable Creating Device Mesh with Physical Axes Splits.


google-cla[bot]: Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/google/jax/pull/20213/checks?check_run_id=22586607283) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Patch:
diff --git a/jax/experimental/mesh_utils.py b/jax/experimental/mesh_utils.py
index 89ef0159d..0d08f62f7 100644
--- a/jax/experimental/mesh_utils.py
+++ b/jax/experimental/mesh_utils.py
@@ -20,10 +20,11 @@ import collections
 from collections.abc import Sequence
 import itertools
 import logging
-from typing import Any, Callable, Optional
+import math
+from typing import Any, Callable, Generator, MutableMapping
 
-import numpy as np
 from jax._src import xla_bridge as xb
+import numpy as np
 
 logger = logging.getLogger(__name__)
 
@@ -36,8 +37,9 @@ _TPU_V4 = 'TPU v4'
 #
 # The trick only works for certain topologies and mesh shapes. Trivial dims of
 # size 1 can be added to the shapes listed, and they are also supported.
-_TRANSPOSE_TRICKS: dict[tuple[int, ...],
-                        dict[tuple[int, ...], tuple[int, ...]]] = {
+_TRANSPOSE_TRICKS: dict[
+    tuple[int, ...], dict[tuple[int, ...], tuple[int, ...]]
+] = {
     (2, 2, 1): {
         (2, 2): (0, 1, 2),
     },
@@ -105,8 +107,11 @@ device_kind_handler_dict: dict[
 
 
 def _create_device_mesh_for_nd_torus(
-    physical_mesh: np.ndarray, mesh_shape: Sequence[int],
-) -> tuple[np.ndarray, list[tuple[int, ...]]]:
+    physical_mesh: np.ndarray,
+    mesh_shape: Sequence[int],
+    *,
+    allow_split_physical_axes: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
   """Assigns logical parallelism axes to physical axes of an N-D torus network.
 
   Given logical parallelism axes with sizes in `mesh_shape` and devices in an
@@ -117,10 +122,9 @@ def _create_device_mesh_for_nd_torus(
   axes of the same size (e.g., a 2D square) rather than multiple physical axes
   of different sizes when possible.
 
-  Note that this routine will never split a physical axis over more than one
-  logical axis (which would reduce total usable bandwidth but may sometimes be
-  desired anyway). As a result, it will error out in cases where this is
-  necessary to produce a valid mapping.
+  If allow_split_physical_axes = False (default), this routine will error out
+  instead of splitting a physical axis over more than one logical axis (which
+  would reduce total usable bandwidth).
 
   Let's use a concrete example to explain the concepts and considerations.
 
@@ -139,12 +143,16 @@ def _create_device_mesh_for_nd_torus(
       physical topology.
     mesh_shape: shape of the logical mesh (size of the various logical
       parallelism axes), with axes ordered by increasing network intensity.
+    allow_split_physical_axes: If True, we would split physical axes if
+      necessary to fit the desired mesh shape.
 
   Returns:
     An np.ndarray of devices in the shape of the logical mesh (mesh_shape), with
       each logical parallelism axis mapped to one or more physical mesh axes.
-    The axis assignment (a list of length num_logical_axes, whose elements
-      are tuples representing physical axis indices).
+    The axis assignment matrix, which is a 2-d array mapping from
+      (physical_axis, logical_axis) to the size assigned, with the invariant
+      np.prod(assignment, axis=1) = physical_mesh_shape, and
+      np.prod(assignment, axis=0) = mesh_shape.
   """
   # Remaining physical axes to be assigned to logical axes.
   assignable_physical_mesh = list(physical_mesh.shape)
@@ -155,14 +163,16 @@ def _create_device_mesh_for_nd_torus(
   # `mesh_shape` is assumed to ordered by lowest network intensity first, so
   # reverse it first.
   for logical_axis_index, logical_axis_size in reversed(
-      list(enumerate(mesh_shape))):
+      list(enumerate(mesh_shape))
+  ):
     # Preferentially map to more physical axes first for higher bandwidth.
     for num_axes in range(3, 0, -1):
       # Try assign to any subset of size num_axes. Generate all candidates.
-      axes = itertools.combinations(assignable_physical_mesh, num_axes)
-      indices = itertools.combinations(
-          range(len(assignable_physical_mesh)), num_axes)
-      for c_axes, c_indices in zip(axes, indices):
+      indices_and_axes = itertools.combinations(
+          enumerate(assignable_physical_mesh), num_axes
+      )
+      for elem in indices_and_axes:
+        c_indices, c_axes = zip(*elem)
         # TODO(zhangqiaorjc): Due to limitations in XLA, 2D collectives only
         # implemented for square 2D plane. Mapping a physical axis to two
         # logical axes might be slower for non-square 2D plane, e.g., map 32 to
@@ -184,20 +194,328 @@ def _create_device_mesh_for_nd_torus(
       # If the num_axes for loop did not break, i.e. none of the candidates work
       # goto here with this while-else construct.
       if logical_axis_size > 1:
-        raise NotImplementedError(
-            'Failed to find assignment for logical_axis_index'
-            f' {logical_axis_index} of size {logical_axis_size} with remaining'
-            f' assignable mesh {assignable_physical_mesh}. The size of each'
-            ' axis in your logical mesh must be equal to the product of'
-            ' some subset of the physical mesh axis sizes. E.g logical mesh (4,'
-            ' 16) is compatible with physical mesh 4x4x4 since 4=4 and 16=4x4.'
-        )
+        if not allow_split_physical_axes:
+          # Although this is now implemented, there are downstream tasks
+          # counting on this being a NotImplementedError.
+          raise NotImplementedError(
+              'Failed to find assignment for logical_axis_index'
+              f' {logical_axis_index} of size {logical_axis_size} with'
+              f' remaining assignable mesh {assignable_physical_mesh}. The size'
+              ' of each axis in your logical mesh must be equal to the product'
+              ' of some subset of the physical mesh axis sizes. E.g. logical'
+              ' mesh (4, 16) is compatible with physical mesh 4x4x4 since 4=4'
+              ' and 16=4x4. If you want to split physical axes, set '
+              ' allow_split_physical_axes to True.'
+          )
+        else:
+          # We will try finding an assignment, even if that means splitting the
+          # physical axes, which requires a more sophisticated implementation.
+          return _create_device_mesh_for_nd_torus_splitting_axes(
+              physical_mesh, mesh_shape
+          )
+
   # Flatten the assignment, e.g., [(), (2,), (0, 1)] -> (2, 0, 1).
   transpose: list[int] = []
-  for x in assignment:
+  assignment_array = np.ones(
+      [len(physical_mesh.shape), len(mesh_shape)], dtype=np.int64
+  )
+  for i, x in enumerate(assignment):
     for y in x:
-      transpose.append(int(y))
-  return physical_mesh.transpose(transpose).reshape(mesh_shape), assignment
+      physical_mesh_axis = int(y)
+      assignment_array[physical_mesh_axis, i] = physical_mesh.shape[
+          physical_mesh_axis
+      ]
+      transpose.append(physical_mesh_axis)
+  return (
+      physical_mesh.transpose(transpose).reshape(mesh_shape),
+      assignment_array,
+  )
+
+
+def _create_device_mesh_for_nd_torus_splitting_axes(
+    physical_mesh: np.ndarray,
+    mesh_shape: Sequence[int],
+) -> tuple[np.ndarray, np.ndarray]:
+  """Assigns logical parallelism axes to physical axes of an N-D torus network.
+
+  This implementation allows creating meshes that requires splitting physical
+  axes, and thus one could produce logical mesh of any shape, as long as the
+  number of devices matches, e.g.,
+
+  - Creating 2x2x4 from 4x4;
+
+  - Creating 2x2x16 from 8x8;
+
+  Args:
+    physical_mesh: a np.ndarray of devices in the shape of the N-D torus
+      physical topology.
+    mesh_shape: shape of the logical mesh (size of the various logical
+      parallelism axes), with axes ordered by increasing network intensity.
+
+  Returns:
+    An np.ndarray of devices in the shape of the logical mesh (mesh_shape), with
+      each logical parallelism axis mapped to one or more physical mesh axes.
+    The axis assignment matrix, which is a 2-d array mapping from
+      (physical_axis, logical_axis) to the size assigned, with the invariant
+      np.prod(assignment, axis=1) = physical_mesh_shape, and
+      np.prod(assignment, axis=0) = mesh_shape.
+  """
+  if np.prod(physical_mesh.shape) != np.prod(mesh_shape):
+    raise ValueError(
+        'The number of devices in physical mesh'
+        f' {physical_mesh.shape} does not match the number of devices'
+        f' in logical mesh {mesh_shape}.'
+    )
+
+  physical_mesh_shape = physical_mesh.shape
+  logical_mesh_shape = tuple(mesh_shape)
+
+  # (Partial) assignment map as an 2-d array [p_axis, l_axis] -> size.
+  assignment = np.ones(
+      [len(physical_mesh_shape), len(logical_mesh_shape)], dtype=np.int64
+  )
+
+  # Process logical axes from highest network intensity to lowest.
+  # `mesh_shape` is assumed to ordered by lowest network intensity first, so
+  # reverse it.
+  for logical_axis, logical_axis_size in reversed(
+      list(enumerate(logical_mesh_shape))
+  ):
+    # Go over all the possible assignment for the logical axis, including the
+    # one that splits multiple physical axes.
+    best_logical_axis_assignment = None
+    for logical_axis_assignment in _enumerate_feasible_logical_axis_assignments(
+        physical_mesh_shape, assignment, logical_axis_size
+    ):
+      # TODO(rosun): Instead of using heuristics, replace this with a proper
+      # scoring function reflecting the underlying hardware properties.
+      if (
+          best_logical_axis_assignment is None
+          or _prefer_first_logical_axis_assignment(
+              logical_axis_assignment,
+              best_logical_axis_assignment,
+              physical_mesh_shape=physical_mesh_shape,
+              assignment=assignment,
+          )
+      ):
+        best_logical_axis_assignment = logical_axis_assignment
+    assignment[:, logical_axis] = best_logical_axis_assignment
+
+  # Read out the assignment.
+  logical_mesh = _generate_logical_mesh(
+      physical_mesh, logical_mesh_shape, assignment
+  )
+
+  return logical_mesh, assignment
+
+
+def _get_prime_factors(x: int) -> list[int]:
+  """Returns a sorted list of prime factors for the given number."""
+  assert x > 0
+  factors = []
+  for p in range(2, math.isqrt(x) + 2):
+    while x % p == 0:
+      factors.append(p)
+      x //= p
+    if x == 1:
+      return factors
+  else:
+    return [x]  # x is a prime number.
+
+
+def _enumerate_feasible_logical_axis_assignments(
+    physical_mesh_shape: Sequence[int],
+    assignment: np.ndarray,
+    logical_axis_size: int,
+) -> Generator[np.ndarray, None, None]:
+  """Yields feasible assignments for a single logical axis.
+
+  For a physical mesh of shape [x_1, ..., x_n], and the product of all previous
+  assignments on each physical axes [y_1, ..., y_n], this function yields all
+  possible assignments for the axis as 1-d arrays [z_1, ..., z_n], so that:
+
+  - prod(z_1, ..., z_n) = logical_axis_size
+
+  - x_i % (z_i * y_i) = 0
+
+  Args:
+    physical_mesh_shape: Physical mesh shape.
+    assignment: Existing assignment matrix.
+    logical_axis_size: Size of the logical axis to assign.
+
+  Yields:
+    All valid assignments for the logical axis. Each assignment is represented
+    as an integer array of length len(physical_mesh_shape).
+  """
+  logical_axis_factors: MutableMapping[int, int] = collections.defaultdict(int)
+  for factor in _get_prime_factors(logical_axis_size):
+    logical_axis_factors[factor] += 1
+
+  available_physical_mesh_shape = np.array(physical_mesh_shape) // np.prod(
+      assignment, axis=-1
+  )
+
+  # To enable efficient enumerations, we first index physical axes by their
+  # prime factors. Since we know the prime factorization of the logical axis
+  # size, we could simply enumerate by picking the correct count for each
+  # prime factor.
+  physical_axes_by_factor: MutableMapping[int, list[int]] = (
+      collections.defaultdict(list)
+  )
+  for physical_axis, physical_axis_size in enumerate(
+      available_physical_mesh_shape
+  ):
+    for factor in _get_prime_factors(physical_axis_size):
+      if factor not in logical_axis_factors:
+        continue
+      physical_axes_by_factor[factor].append(physical_axis)
+
+  factors = []
+  assignments_by_factor = []
+  for factor, multiplicity in logical_axis_factors.items():
+    factors.append(factor)
+    assignments_by_factor.append(
+        set(
+            itertools.combinations(
+                physical_axes_by_factor[factor], multiplicity
+            )
+        )
+    )
+
+  for axis_assignment in itertools.product(*assignments_by_factor):
+    result = np.ones([len(physical_mesh_shape)], dtype=np.int64)
+    for factor_index, per_factor_assignment in enumerate(axis_assignment):
+      for physical_axis in per_factor_assignment:
+        result[physical_axis] *= factors[factor_index]
+    yield result
+
+
+def _prefer_first_logical_axis_assignment(
+    x: np.ndarray,
+    y: np.ndarray,
+    *,
+    physical_mesh_shape: Sequence[int],
+    assignment: np.ndarray,
+) -> bool:
+  """Returns True if the first axis assignment is preferred over the second.
+
+  For now, this is implemented with some very simple heuristics. However,
+  it is possible to introduce e.g., a value function here based on a more
+  precise model of the underlying hardware.
+
+  TODO(rosun): Use a proxy of network capacity to select the partitions.
+
+  Args:
+    x: Logical axis assignment as [len(physical_mesh_shape)] array.
+    y: Logical axis assignment as [len(physical_mesh_shape)] array.
+    physical_mesh_shape: Physical mesh shape.
+    assignment: Assignment matrix.
+
+  Returns:
+    True if x is preferred over y.
+  """
+  # Prefer occupying complete physical axes. I don't have a good reason for
+  # this, except that it is compatible with the existing behavior.
+  #
+  # E.g., on 4 x 4 x 8, [4, 4, -] will be preferred over [4, -, 4], and then
+  # over [2, 2, 4].
+  x_whole_axis_size = np.prod(
+      [s for i, s in enumerate(x) if s == physical_mesh_shape[i]]
+  )
+  y_whole_axis_size = np.prod(
+      [s for i, s in enumerate(y) if s == physical_mesh_shape[i]]
+  )
+
+  if x_whole_axis_size != y_whole_axis_size:
+    return x_whole_axis_size > y_whole_axis_size
+
+  # Prefer occupying more whole physical axes for better bandwidth.
+  #
+  # This is consistent with existing logic, i.e., 2 x 2 is preferred over 4.
+  x_num_whole_axes = len(
+      [1 for i, s in enumerate(x) if s == physical_mesh_shape[i] and s > 1]
+  )
+  y_num_whole_axes = len(
+      [1 for i, s in enumerate(y) if s == physical_mesh_shape[i] and s > 1]
+  )
+
+  if x_num_whole_axes != y_num_whole_axes:
+    return x_num_whole_axes > y_num_whole_axes
+
+  # Prefer taking physical axes that are not taken by logical axes of higher
+  # network intensity. E.g., for a 4 x 4 x 4, suppose that the previous
+  # assignments are 1 x 2 x 4, and we want to place a new logical axis of size
+  # 2, we will go for [2, 1, 1] instead of [1, 2, 1], as the latter choice will
+  # tap into bandwidth already taken by the higher intensity axis.
+  assigned_physical_mesh_shape = np.prod(assignment, axis=-1)
+
+  x_non_overlapping_axis_size = np.prod(
+      [s for i, s in enumerate(x) if assigned_physical_mesh_shape[i] > 1]
+  )
+  y_non_overlapping_axis_size = np.prod(
+      [s for i, s in enumerate(y) if assigned_physical_mesh_shape[i] > 1]
+  )
+
+  if x_non_overlapping_axis_size != y_non_overlapping_axis_size:
+    return x_non_overlapping_axis_size > y_non_overlapping_axis_size
+
+  # Otherwise sort by reverse lexical graphical order, to be consistent with
+  # existing behavior.
+  return tuple(x) > tuple(y)
+
+
+def _generate_logical_mesh(
+    physical_mesh: np.ndarray,
+    logical_mesh_shape: Sequence[int],
+    assignment: np.ndarray,
+) -> np.ndarray:
+  """Compute the logical mesh from assignment map.
+
+  Args:
+    physical_mesh: Physical device mesh.
+    logical_mesh_shape: Logical mesh shape.
+    assignment: 2-d assignment matrix shape [physical_dims, logical_dims].
+
+  Returns:
+    Logical mesh reshaped from physical mesh.
+  """
+  physical_indices = np.broadcast_to(
+      np.expand_dims(
+          np.arange(len(physical_mesh.shape), dtype=np.int64), axis=-1
+      ),
+      assignment.shape,
+  ).reshape([-1])
+
+  logical_indices = np.broadcast_to(
+      np.expand_dims(
+          np.arange(len(logical_mesh_shape), dtype=np.int64), axis=0
+      ),
+      assignment.shape,
+  ).reshape([-1])
+
+  # Axes of logical mesh is ordered by (physical_axis, logical_axis).
+  #
+  # Note that we sort for each physical_axis the logical_axis, so that higher
+  # intensity logical axes are replicated at inner (minor) dimensions.
+  #
+  # E.g., if a dimension size is 12 = 3x4, where 3 is higher intensity and 4
+  # is lower, we want to reshape so that it becomes 12 = 4x3. Imagine in the
+  # 1-d case, this will allow more connections between the higher intensity
+  # axes.
+  logical_mesh = np.reshape(physical_mesh, assignment.reshape([-1]))
+
+  # We will then group by l_axis as this is what is expected from output.
+  _, _, transpose_axes = zip(
+      *sorted(
+          zip(logical_indices, physical_indices, range(len(logical_indices)))
+      )
+  )
+  logical_mesh = np.transpose(logical_mesh, transpose_axes)
+
+  # Reshape to add the trivial dimensions back.
+  logical_mesh = np.reshape(logical_mesh, logical_mesh_shape)
+
+  return logical_mesh
 
 
 def _bounds_from_last_device(last_device) -> Sequence[int]:
@@ -245,14 +563,16 @@ def _get_physical_tpu_mesh(jax_devices: Sequence[Any]) -> np.ndarray:
 
 
 # jekbradbury's famous trick for creating contiguous submeshes (where available)
-def _transpose_trick(physical_mesh: np.ndarray,
-                     mesh_shape: Sequence[int]) -> np.ndarray:
+def _transpose_trick(
+    physical_mesh: np.ndarray, mesh_shape: Sequence[int]
+) -> np.ndarray:
   mesh_shape = tuple(mesh_shape)
   topology = physical_mesh.shape
   if topology not in _TRANSPOSE_TRICKS:
     raise ValueError(
-        f"create_device_mesh cannot create contiguous submeshes for "
-        f"physical mesh topology {topology}")
+        'create_device_mesh cannot create contiguous submeshes for '
+        f'physical mesh topology {topology}'
+    )
 
   mesh_shape_no_trivial_dims: tuple[int, ...] = ()
   for dim_size in mesh_shape:
@@ -261,18 +581,23 @@ def _transpose_trick(physical_mesh: np.ndarray,
 
   if mesh_shape_no_trivial_dims not in _TRANSPOSE_TRICKS[topology]:
     raise ValueError(
-        f"create_device_mesh cannot create contiguous submeshes for "
-        f"mesh_shape {mesh_shape} and physical mesh topology {topology}. "
-        f"Available mesh_shapes: {list(_TRANSPOSE_TRICKS[topology].keys())}")
+        'create_device_mesh cannot create contiguous submeshes for '
+        f'mesh_shape {mesh_shape} and physical mesh topology {topology}. '
+        f'Available mesh_shapes: {list(_TRANSPOSE_TRICKS[topology].keys())}'
+    )
 
   return physical_mesh.transpose(
-      *_TRANSPOSE_TRICKS[topology][mesh_shape_no_trivial_dims])
+      *_TRANSPOSE_TRICKS[topology][mesh_shape_no_trivial_dims]
+  )
 
 
 def create_device_mesh(
     mesh_shape: Sequence[int],
-    devices: Sequence[Any] | None = None, *,
-    contiguous_submeshes: bool = False) -> np.ndarray:
+    devices: Sequence[Any] | None = None,
+    *,
+    contiguous_submeshes: bool = False,
+    allow_split_physical_axes: bool = False,
+) -> np.ndarray:
   """Creates a performant device mesh for jax.sharding.Mesh.
 
   Args:
@@ -287,6 +612,8 @@ def create_device_mesh(
       setting was sometimes necessary before the introduction of jax.Array to
       ensure non-ragged local arrays; if using jax.Arrays, it's better to keep
       this set to False.
+    allow_split_physical_axes: If True, we will split physical axes if necessary
+      to produce the desired device mesh.
 
   Raises:
     ValueError: if the number of devices doesn't equal the product of
@@ -299,8 +626,10 @@ def create_device_mesh(
   if devices is None:
     devices = xb.devices()
   if np.prod(mesh_shape) != len(devices):
-    raise ValueError(f'Number of devices {len(devices)} must equal the product '
-                     f'of mesh_shape {mesh_shape}')
+    raise ValueError(
+        f'Number of devices {len(devices)} must equal the product '
+        f'of mesh_shape {mesh_shape}'
+    )
   last_device = devices[-1]
 
   handler = device_kind_handler_dict.get(last_device.device_kind, None)
@@ -315,18 +644,25 @@ def create_device_mesh(
     physical_mesh = _get_physical_tpu_mesh(devices)
     if contiguous_submeshes:
       physical_mesh = _transpose_trick(physical_mesh, mesh_shape)
-    device_mesh, _ = _create_device_mesh_for_nd_torus(physical_mesh, mesh_shape)
+    device_mesh, _ = _create_device_mesh_for_nd_torus(
+        physical_mesh,
+        mesh_shape,
+        allow_split_physical_axes=allow_split_physical_axes,
+    )
     return device_mesh
   else:
     device_mesh = np.asarray(devices).reshape(mesh_shape)
     return device_mesh
 
+
 def create_hybrid_device_mesh(
     mesh_shape: Sequence[int],
     dcn_mesh_shape: Sequence[int],
-    devices: Sequence[Any] | None = None, *,
+    devices: Sequence[Any] | None = None,
+    *,
     process_is_granule: bool = False,
     should_sort_granules_by_key: bool = True,
+    allow_split_physical_axes: bool = False,
 ) -> np.ndarray:
   """Creates a device mesh for hybrid (e.g., ICI and DCN) parallelism.
 
@@ -345,6 +681,8 @@ def create_hybrid_device_mesh(
     should_sort_granules_by_key: Whether device granules should be sorted by the
       granule key, either slice or process index, depending on
       process_is_granule.
+    allow_split_physical_axes: If True, we will split physical axes if necessary
+      to produce the desired device mesh.
 
   Raises:
     ValueError: if the number of slices to which the `devices` belong doesn't
@@ -365,16 +703,25 @@ def create_hybrid_device_mesh(
   granules = (
       [granule_dict[key] for key in sorted(granule_dict.keys())]
       if should_sort_granules_by_key
-      else granule_dict.values())
+      else granule_dict.values()
+  )
   if np.prod(dcn_mesh_shape) != len(granules):
     raise ValueError(
         f'Number of slices {len(granules)} must equal the product of '
-        f'dcn_mesh_shape {dcn_mesh_shape}')
-  per_granule_meshes = [create_device_mesh(mesh_shape, granule)
-                        for granule in granules]
+        f'dcn_mesh_shape {dcn_mesh_shape}'
+    )
+  per_granule_meshes = [
+      create_device_mesh(
+          mesh_shape,
+          granule,
+          allow_split_physical_axes=allow_split_physical_axes,
+      )
+      for granule in granules
+  ]
   # TODO(jekbradbury): handle non-uniform DCN topologies
   granule_mesh = np.arange(len(granules)).reshape(dcn_mesh_shape)
-  blocks = np.vectorize(
-    lambda i: per_granule_meshes[i], otypes=[object])(granule_mesh)
+  blocks = np.vectorize(lambda i: per_granule_meshes[i], otypes=[object])(
+      granule_mesh
+  )
   device_mesh = np.block(blocks.tolist())
   return device_mesh


Test patch:
diff --git a/tests/mesh_utils_test.py b/tests/mesh_utils_test.py
index 0ed5ce4c5..420eb9235 100644
--- a/tests/mesh_utils_test.py
+++ b/tests/mesh_utils_test.py
@@ -18,14 +18,13 @@ import collections
 from collections.abc import Sequence
 import dataclasses
 
-import numpy as np
-
 from absl import logging
 from absl.testing import absltest
 from absl.testing import parameterized
-from jax.experimental import mesh_utils
-from jax.sharding import Mesh
 from jax._src import test_util
+from jax.experimental import mesh_utils
+from jax.sharding import Mesh  # pylint: disable=g-importing-member
+import numpy as np
 
 
 @dataclasses.dataclass(frozen=True)
@@ -183,15 +182,81 @@ class MeshUtilsTest(test_util.JaxTestCase):
       ('4x8x8', mock_4x8x8_devices, [1, 32, 8], [(), (0, 2), (1,)]),
       ('8x8x8', mock_8x8x8_devices, [1, 64, 8], [(), (1, 2), (0,)]),
       ('8x8x16', mock_8x8x16_devices, [1, 64, 16], [(), (0, 1), (2,)]),
-      ('8x8', mock_8x8_devices, [8, 8], [(1,), (0, 2)])
+      ('8x8', mock_8x8_devices, [8, 8], [(1,), (0, 2)]),
   )
-  def test_create_device_mesh_for_nd_torus(self, devices, mesh_shape,
-                                           expected_assignment):
+  def test_create_device_mesh_for_nd_torus(
+      self, devices, mesh_shape, expected_assignment
+  ):
     jax_devices = devices(True)
     physical_mesh = mesh_utils._get_physical_tpu_mesh(jax_devices)
     _, assignment = mesh_utils._create_device_mesh_for_nd_torus(
-        physical_mesh, mesh_shape)
-    self.assertEqual(assignment, expected_assignment)
+        physical_mesh, mesh_shape
+    )
+
+    # The expected assignment is specified as a list, where each element is a
+    # sequence of physical axis assigned. We convert this into assignment
+    # matrix.
+    expected_assignment_matrix = np.ones(
+        [physical_mesh.ndim, len(mesh_shape)], dtype=np.int64
+    )
+    for logical_axis, axis_assignment in enumerate(expected_assignment):
+      for physical_axis in axis_assignment:
+        expected_assignment_matrix[physical_axis, logical_axis] = (
+            physical_mesh.shape[physical_axis]
+        )
+    self.assertArraysEqual(assignment, expected_assignment_matrix)
+
+  @parameterized.named_parameters(
+      ('2x2x1', mock_2x2x1_devices, [1, 1, 4], [(), (), (0, 1, 2)]),
+      ('2x2x4', mock_2x2x4_devices, [1, 4, 4], [(), (2,), (0, 1)]),
+      ('4x4x4', mock_4x4x4_devices, [1, 16, 4], [(), (1, 2), (0,)]),
+      ('4x4x8a', mock_4x4x8_devices, [1, 16, 8], [(), (0, 1), (2,)]),
+      ('4x4x8b', mock_4x4x8_devices, [1, 8, 16], [(), (2,), (0, 1)]),
+      ('4x4x8c', mock_4x4x8_devices, [16, 8, 1], [(0, 1), (2,), ()]),
+      ('4x8x8', mock_4x8x8_devices, [1, 32, 8], [(), (0, 2), (1,)]),
+      ('8x8x8', mock_8x8x8_devices, [1, 64, 8], [(), (1, 2), (0,)]),
+      ('8x8x16', mock_8x8x16_devices, [1, 64, 16], [(), (0, 1), (2,)]),
+      ('8x8', mock_8x8_devices, [8, 8], [(1,), (0, 2)]),
+  )
+  def test_create_device_mesh_for_nd_torus_split_axes_backward_compatible(
+      self, devices, mesh_shape, expected_assignment
+  ):
+    jax_devices = devices(True)
+    physical_mesh = mesh_utils._get_physical_tpu_mesh(jax_devices)
+    _, assignment = mesh_utils._create_device_mesh_for_nd_torus_splitting_axes(
+        physical_mesh, mesh_shape
+    )
+
+    # The expected assignment is specified as a list, where each element is a
+    # sequence of physical axis assigned. We convert this into assignment
+    # matrix.
+    expected_assignment_matrix = np.ones(
+        [physical_mesh.ndim, len(mesh_shape)], dtype=np.int64
+    )
+    for logical_axis, axis_assignment in enumerate(expected_assignment):
+      for physical_axis in axis_assignment:
+        expected_assignment_matrix[physical_axis, logical_axis] = (
+            physical_mesh.shape[physical_axis]
+        )
+    self.assertArraysEqual(assignment, expected_assignment_matrix)
+
+  @parameterized.named_parameters(
+      ('4x4x4a', mock_4x4x4_devices, [2, 1, 32]),
+      ('4x4x4b', mock_4x4x4_devices, [8, 8, 1]),
+      ('4x4x8a', mock_4x4x8_devices, [2, 2, 8, 4]),
+      ('4x4x8b', mock_4x4x8_devices, [2, 4, 1, 16]),
+      ('4x8x8', mock_4x8x8_devices, [1, 128, 2]),
+      ('8x8', mock_8x8_devices, [2, 1, 32, 1]),
+  )
+  def test_create_device_mesh_for_nd_torus_split_axes_can_handle_axes_split(
+      self, devices, mesh_shape
+  ):
+    jax_devices = devices(True)
+    physical_mesh = mesh_utils._get_physical_tpu_mesh(jax_devices)
+    logical_mesh, _ = mesh_utils._create_device_mesh_for_nd_torus(
+        physical_mesh, mesh_shape, allow_split_physical_axes=True
+    )
+    self.assertEqual(logical_mesh.shape, tuple(mesh_shape))
 
   @parameterized.named_parameters(
       ('2X4x4x4a', (1, 16, 4), (2, 1, 1)),
@@ -324,7 +389,160 @@ class MeshUtilsTest(test_util.JaxTestCase):
         "(1, 128, 2) and physical mesh topology (4, 8, 8). "
         'Available mesh_shapes: [(64, 4), (4, 64)]'):
       mesh_utils.create_device_mesh(
-          mesh_shape, devices=devices, contiguous_submeshes=True)
+          mesh_shape, devices=devices, contiguous_submeshes=True
+      )
+
+
+def int64_array(x) -> np.ndarray:
+  return np.array(x, dtype=np.int64)
+
+
+def get_int_mesh(shape: Sequence[int]) -> np.ndarray:
+  return np.arange(np.prod(shape), dtype=np.int64).reshape(shape)
+
+
+class SplitAxesDeviceMeshCreationTest(test_util.JaxTestCase):
+
+  def test_get_prime_factors(self):
+    self.assertEqual(mesh_utils._get_prime_factors(1), [])  # 1 has no factor.
+    self.assertEqual(mesh_utils._get_prime_factors(2), [2])
+    self.assertEqual(mesh_utils._get_prime_factors(4), [2, 2])
+    self.assertEqual(mesh_utils._get_prime_factors(8), [2, 2, 2])
+    self.assertEqual(mesh_utils._get_prime_factors(6), [2, 3])
+    self.assertEqual(mesh_utils._get_prime_factors(16), [2, 2, 2, 2])
+    self.assertEqual(mesh_utils._get_prime_factors(12), [2, 2, 3])
+    self.assertEqual(mesh_utils._get_prime_factors(121), [11, 11])  # square
+    self.assertEqual(mesh_utils._get_prime_factors(43), [43])  # prime
+
+  @parameterized.named_parameters(
+      (
+          '2x2x1',
+          [2, 2, 1],
+          [1, 2, 1],
+          4,
+          [],  # infeasible
+      ),
+      (
+          '12x4x4',
+          [12, 4, 4],
+          [2, 2, 1],
+          6,
+          [[6, 1, 1], [3, 2, 1], [3, 1, 2]],
+      ),
+      (
+          '4x4x8',
+          [4, 4, 8],
+          [2, 2, 2],
+          4,
+          [[2, 2, 1], [2, 1, 2], [1, 2, 2], [1, 1, 4]],
+      ),
+  )
+  def test_enumerate_feasible_axis_assignments(
+      self,
+      physical_mesh_shape,
+      assigned_physical_mesh_shape,
+      logical_axis_size,
+      expected_assignments,
+  ):
+    assignment = int64_array([list(assigned_physical_mesh_shape)]).T
+    self.assertArraysEqual(
+        list(
+            mesh_utils._enumerate_feasible_logical_axis_assignments(
+                physical_mesh_shape,
+                assignment,
+                logical_axis_size=logical_axis_size,
+            )
+        ),
+        [int64_array(a) for a in expected_assignments],
+    )
+
+  @parameterized.named_parameters(
+      (
+          '2x2x1',
+          [2, 2, 1],
+          [1, 2, 2, 1],
+          [
+              [1, 2, 1, 1],
+              [1, 1, 2, 1],
+              [1, 1, 1, 1],
+          ],
+      ),
+      (
+          '4x4x4',
+          [4, 4, 4],
+          [2, 1, 32],
+          [
+              [1, 1, 4],
+              [2, 1, 2],
+              [1, 1, 4],
+          ],
+      ),
+      (
+          '12x4x8',
+          [12, 4, 8],
+          [2, 8, 24],
+          [
+              [2, 2, 3],
+              [1, 2, 4],
+              [1, 2, 2],
+          ],
+      ),
+  )
+  def test_generate_logical_mesh(
+      self,
+      physical_mesh_shape,
+      logical_mesh_shape,
+      assignment,
+  ):
+    assignment = np.array(assignment, dtype=np.int64)
+    physical_mesh = get_int_mesh(physical_mesh_shape)
+    logical_mesh = mesh_utils._generate_logical_mesh(
+        physical_mesh, logical_mesh_shape, assignment
+    )
+    self.assertEqual(logical_mesh.shape, tuple(logical_mesh_shape))
+    # We check that the logical mesh is assigned correctly using the following
+    # consistency check, which transforms the logical mesh back to physical
+    # mesh.
+    transpose = (
+        np.arange(assignment.size).reshape(assignment.shape).T.reshape([-1])
+    )
+    self.assertArraysEqual(
+        physical_mesh.reshape([-1]),
+        logical_mesh.reshape(np.reshape(assignment.T, [-1]))
+        .transpose(transpose)
+        .reshape([-1]),
+    )
+
+  def test_prefer_assignment_whole_axis_size(self):
+    self.assertTrue(
+        mesh_utils._prefer_first_logical_axis_assignment(
+            int64_array([1, 2, 1]),
+            int64_array([1, 1, 2]),
+            physical_mesh_shape=[2, 2, 4],
+            assignment=int64_array([[1, 1, 1]]).T,
+        )
+    )
+
+  def test_prefer_assignment_more_whole_axes(self):
+    # This entails the original implementation already.
+    self.assertTrue(
+        mesh_utils._prefer_first_logical_axis_assignment(
+            int64_array([4, 4, 1]),
+            int64_array([1, 1, 16]),
+            physical_mesh_shape=[4, 4, 16],
+            assignment=int64_array([[1, 1, 1]]).T,
+        )
+    )
+
+  def test_prefer_assignment_avoid_already_assigned(self):
+    self.assertTrue(
+        mesh_utils._prefer_first_logical_axis_assignment(
+            int64_array([2, 1]),
+            int64_array([1, 2]),
+            physical_mesh_shape=[2, 4],
+            assignment=int64_array([[1, 2]]).T,
+        )
+    )
 
 
 if __name__ == '__main__':