jax-ml/jax · critical · ValueError

When passing host local inputs to pjit, devices connected to

Error message

When passing host local inputs to pjit, devices connected to a single host must form a contiguous subcube of the global device mesh

What it means

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.

Source

Thrown at jax/_src/mesh.py:104

  is_local_device = np.vectorize(
      lambda d: d.process_index == process_index, otypes=[bool])(global_mesh.devices)
  subcube_indices = []
  # We take the smallest slice of each dimension that doesn't skip any local device.
  for axis in range(global_mesh.devices.ndim):
    other_axes = tuple_delete(tuple(range(global_mesh.devices.ndim)), axis)
    # NOTE: This re-reduces over many axes multiple times, so we could definitely
    #       optimize it, but I hope it won't be a bottleneck anytime soon.
    local_slices = is_local_device.any(other_axes, keepdims=False)
    nonzero_indices = np.flatnonzero(local_slices)
    start, end = int(np.min(nonzero_indices)), int(np.max(nonzero_indices))
    subcube_indices.append(slice(start, end + 1))
  subcube_indices_tuple = tuple(subcube_indices)
  # We only end up with all conditions being true if the local devices formed a
  # subcube of the full array. This is because we were biased towards taking a
  # "hull" spanned by the devices, and in case the local devices don't form a
  # subcube that hull will contain non-local devices.
  if not is_local_device[subcube_indices_tuple].all():
    raise ValueError(
        "When passing host local inputs to pjit, devices connected to a single"
        " host must form a contiguous subcube of the global device mesh"
    )
  return Mesh(global_mesh.devices[subcube_indices_tuple], global_mesh.axis_names)


class AxisType(enum.Enum):
  Auto = enum.auto()
  Explicit = enum.auto()
  Manual = enum.auto()

  def __repr__(self):
    return self.name

def _normalize_axis_types(axis_names, axis_types, name, default_axis_type):
  axis_types = ((default_axis_type,) * len(axis_names)
                if axis_types is None else axis_types)
  if not isinstance(axis_types, tuple):

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
devices = sorted(jax.devices(), key=lambda d: d.id)  # may interleave hosts
mesh = jax.sharding.Mesh(np.array(devices).reshape(4, 8), ('x', 'y'))

# after
from jax.experimental.mesh_utils import create_device_mesh
mesh = jax.sharding.Mesh(create_device_mesh((4, 8)), ('x', 'y'))

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np, jax
devs = np.array(jax.devices())
# ensure each host's devices are contiguous along some subcube: simplest check:
host_ids = [d.host_id for d in devs.flat]
assert host_ids == sorted(host_ids), 'devices not grouped by host'

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e13bbe5337ddb02b. Report an issue: GitHub.