{"record":{"id":"e13bbe5337ddb02b","repo":"jax-ml/jax","slug":"when-passing-host-local-inputs-to-pjit-devices-co","errorCode":null,"errorMessage":"When passing host local inputs to pjit, devices connected to a single host must form a contiguous subcube of the global device mesh","messagePattern":"When passing host local inputs to pjit, devices connected to a single host must form a contiguous subcube of the global device mesh","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"critical","filePath":"jax/_src/mesh.py","lineNumber":104,"sourceCode":"  is_local_device = np.vectorize(\n      lambda d: d.process_index == process_index, otypes=[bool])(global_mesh.devices)\n  subcube_indices = []\n  # We take the smallest slice of each dimension that doesn't skip any local device.\n  for axis in range(global_mesh.devices.ndim):\n    other_axes = tuple_delete(tuple(range(global_mesh.devices.ndim)), axis)\n    # NOTE: This re-reduces over many axes multiple times, so we could definitely\n    #       optimize it, but I hope it won't be a bottleneck anytime soon.\n    local_slices = is_local_device.any(other_axes, keepdims=False)\n    nonzero_indices = np.flatnonzero(local_slices)\n    start, end = int(np.min(nonzero_indices)), int(np.max(nonzero_indices))\n    subcube_indices.append(slice(start, end + 1))\n  subcube_indices_tuple = tuple(subcube_indices)\n  # We only end up with all conditions being true if the local devices formed a\n  # subcube of the full array. This is because we were biased towards taking a\n  # \"hull\" spanned by the devices, and in case the local devices don't form a\n  # subcube that hull will contain non-local devices.\n  if not is_local_device[subcube_indices_tuple].all():\n    raise ValueError(\n        \"When passing host local inputs to pjit, devices connected to a single\"\n        \" host must form a contiguous subcube of the global device mesh\"\n    )\n  return Mesh(global_mesh.devices[subcube_indices_tuple], global_mesh.axis_names)\n\n\nclass AxisType(enum.Enum):\n  Auto = enum.auto()\n  Explicit = enum.auto()\n  Manual = enum.auto()\n\n  def __repr__(self):\n    return self.name\n\ndef _normalize_axis_types(axis_names, axis_types, name, default_axis_type):\n  axis_types = ((default_axis_type,) * len(axis_names)\n                if axis_types is None else axis_types)\n  if not isinstance(axis_types, tuple):","sourceCodeStart":86,"sourceCodeEnd":122,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L86-L122","documentation":"When building a per-host local mesh for host-local (INPUT_PER_DEVICE-style) pjit inputs, JAX computes the hypercube hull of the host's devices within the global mesh. If the hull contains non-local devices, the host's devices are not a contiguous subcube, and _get_local_mesh raises ValueError because host-local arrays cannot be laid out on such a set.","triggerScenarios":"Constructing a global Mesh where the devices assigned to one host form a non-contiguous pattern (e.g. interleaved or L-shaped) and then passing host-local inputs to pjit/jit; common when device_order or a custom device list scrambles host locality.","commonSituations":"Custom device meshes built with jax.sharding.Mesh over an explicitly ordered device list that doesn't group each host's devices contiguously; multi-host TPU/pmap-to-pjit migration; using devices sorted by global device id rather than by host.","solutions":["Reorder devices so each host's devices form a contiguous subcube of the mesh before constructing the Mesh","Use jax.experimental.mesh_utils.create_device_mesh / create_hybrid_device_mesh which produce host-contiguous orderings","Avoid hand-rolled device orderings; derive the mesh from jax.devices() with locality-aware helpers","If interleaving is intentional, don't use host-local input layout — pass global arrays with NamedSharding"],"exampleFix":"# before\ndevices = sorted(jax.devices(), key=lambda d: d.id)  # may interleave hosts\nmesh = jax.sharding.Mesh(np.array(devices).reshape(4, 8), ('x', 'y'))\n\n# after\nfrom jax.experimental.mesh_utils import create_device_mesh\nmesh = jax.sharding.Mesh(create_device_mesh((4, 8)), ('x', 'y'))","handlingStrategy":"validation","validationCode":"import numpy as np, jax\ndevs = np.array(jax.devices())\n# ensure each host's devices are contiguous along some subcube: simplest check:\nhost_ids = [d.host_id for d in devs.flat]\nassert host_ids == sorted(host_ids), 'devices not grouped by host'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always build meshes with create_device_mesh/create_hybrid_device_mesh","Never hand-sort devices by id across hosts","Group each host's devices contiguously in the device array"],"tags":["jax","mesh","multi-host","tpu","device-topology"],"backgroundTag":"non-contiguous-device-mesh","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}