code: def _get_physical_tpu_mesh(jax_devices: Sequence[Any]) -> np.ndarray:
      r"""Rearrange TPU devices in a slice into a physical mesh.
  
      Args:
        jax_devices: A list of JAX devices in a TPU slice in process-tiled z, y, x,
          core order, e.g. from jax.devices().
  
      Returns:
        A np.ndarray of JAX devices with shape [global_x, global_y, global_z]. On
          v2 and v3, global_z is instead cores_per_chip (i.e., 2).
      """
      device_kind = jax_devices[0].device_kind
      device_coords = [d.coords for d in jax_devices]
      dims = tuple(d + 1 for d in max(device_coords))
      assert len(dims) == 3, dims
      if device_kind in (_TPU_V2, _TPU_V3):
          cores_per_chip = max(d.core_on_chip for d in jax_devices) + 1
          out = np.empty(dims[:2] + (cores_per_chip,), dtype=object)
          for coords, d in zip(device_coords, jax_devices):
              assert coords[2] == 0, d
              out[coords[0], coords[1], d.core_on_chip] = d
      else:
          out = np.empty(dims, dtype=object)
          for coords, d in zip(device_coords, jax_devices):
              if d.core_on_chip != 0:
                  raise AssertionError(
                      "Creating meshes for TPU >v3 requires one device per chip"
                      f' ("megacore" mode). Got device id {d.core_on_chip} for a device'
                      f" of kind {device_kind}: {d}."
                  )
              out[coords[0], coords[1], coords[2]] = d
      return out
  
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 4 for calling a global function with example usage: _get_physical_tpu_mesh(jax_devices) and returns: out