code: def create_device_mesh(
      mesh_shape: Sequence[int],
      devices: Sequence[Any] | None = None,
      *,
      contiguous_submeshes: bool = False,
  ) -> np.ndarray:
      """Creates a performant device mesh for jax.sharding.Mesh.
  
      Args:
        mesh_shape: shape of logical mesh, ordered by increasing network-intensity
          e.g. [replica, data, mdl] where mdl has the most network communication
          requirements.
        devices: optionally, the devices to construct a mesh for. Defaults to
          jax.devices().
        contiguous_submeshes: if True, this function will attempt to create a mesh
          where each process's local devices form a contiguous submesh. A ValueError
          will be raised if this function can't produce a suitable mesh. This
          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.
  
      Raises:
        ValueError: if the number of devices doesn't equal the product of
          `mesh_shape`.
  
      Returns:
        A np.ndarray of JAX devices with mesh_shape as its shape that can be fed
        into jax.sharding.Mesh with good collective performance.
      """
      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}"
          )
      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,
          )
          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
              )
          device_mesh, assignment = _create_device_mesh_for_nd_torus(
              physical_mesh, mesh_shape
          )
          logger.info(
              "_create_device_mesh_for_nd_torus assignment: %s",
              assignment,
          )
          return device_mesh
      else:
          device_mesh = np.asarray(devices).reshape(mesh_shape)
          return device_mesh
  
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 6 for calling a global function with example usage: create_device_mesh(mesh_shape, devices) and returns: device_mesh, device_mesh, result