code: 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}"
          )
  
      mesh_shape_no_trivial_dims: tuple[int, ...] = ()
      for dim_size in mesh_shape:
          if dim_size != 1:
              mesh_shape_no_trivial_dims += (dim_size,)
  
      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())}"
          )
  
      return physical_mesh.transpose(
          *_TRANSPOSE_TRICKS[topology][mesh_shape_no_trivial_dims]
      )
  
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 5 for calling a global function with example usage: _transpose_trick(physical_mesh, mesh_shape) and returns: physical_mesh