code: def _create_device_mesh_for_nd_torus(
      physical_mesh: np.ndarray,
      mesh_shape: Sequence[int],
  ) -> tuple[np.ndarray, list[tuple[int, ...]]]:
      """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
      N-dimensional torus network represented by `physical_mesh`, maps each logical
      axis to one or more physical axes. Prefer to map more-performance-sensitive
      logical axes to larger numbers of physical axes to maximize the bandwidth
      available to them. Also prefer to assign logical axes to multiple physical
      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.
  
      Let's use a concrete example to explain the concepts and considerations.
  
      As an example, suppose the logical mesh is [data, model], for data and model
      parallelism respectively. Also suppose that data parallelism is less
      performance sensitive than model parallelism. Consider a 3D TPU pod slice of
      shape 4x4x16, represented by a physical mesh of shape (4, 4, 16).
  
      A TPU pod slice has equal bandwidth along all axes with wraparound links, but
      a 2D plane of size 4x4 may have faster XLA collective implementations than a
      non-square plane or a 1D subgroup. If the mesh_shape is [16, 16], we may want
      the more performance sensitive `model` axis to be mapped to the 4x4 XY plane.
  
      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 (a list of length num_logical_axes, whose elements
          are tuples representing physical axis indices).
      """
      # Remaining physical axes to be assigned to logical axes.
      assignable_physical_mesh = list(physical_mesh.shape)
      # Map each logical axis to a subset of physical axes.
      assignment: list[tuple[int, ...]] = [() for _ in mesh_shape]
  
      # Assign logical axes from highest network intensity to lowest.
      # `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))
      ):
          # 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):
                  # 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
                  # 4x8 or a single axis. If XLA 2D collectives support non-square plane
                  # soon, we can continue to preferentially map to 2D plane in general,
                  # otherwise, we should treat non-square 2D plane and 1D submesh equally.
                  if np.prod(c_axes) == logical_axis_size:
                      assignment[logical_axis_index] = c_indices
                      # Zero the assigned physical axes.
                      assignable_physical_mesh = [
                          0 if i in c_indices else v
                          for i, v in enumerate(
                              assignable_physical_mesh
                          )
                      ]
                      break
              if assignment[logical_axis_index]:
                  # We already found an assignment from one candidate above.
                  break
          else:
              # 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."
                  )
      # Flatten the assignment, e.g., [(), (2,), (0, 1)] -> (2, 0, 1).
      transpose: list[int] = []
      for x in assignment:
          for y in x:
              transpose.append(int(y))
      return (
          physical_mesh.transpose(transpose).reshape(mesh_shape),
          assignment,
      )
  
file_path: jax/jax/experimental/mesh_utils.py
imports: from __future__ import annotations
  import collections
  from collections.abc import Sequence
  import itertools
  import logging
  from typing import Any, Callable, Optional
  import numpy
  from jax._src import xla_bridge
label_desc: Use image node_id 2 for calling a global function with example usage: _create_device_mesh_for_nd_torus(physical_mesh, mesh_shape) and returns: physical_mesh,  assignment