jax-ml/jax · error · ValueError

Number of device ids must match the number of mesh axes, but

Error message

Number of device ids must match the number of mesh axes, but got {len(device_ids)} ids for a {len(mesh_strides)}D mesh.

What it means

When device_id_type is MESH, the number of device ids must equal the mesh dimensionality; got fewer/more ids than mesh axes.

Source

Thrown at jax/_src/pallas/primitives.py:1282

  """
  non_mesh_axes = {}
  if isinstance(device_id, dict):
    if device_id_type is not DeviceIdType.MESH:
      raise ValueError(
          "`device_id_type` must be MESH if `device_id` is a dict,"
          f" got: {device_id_type = }."
      )
    device_id, non_mesh_axes = _device_id_dict_to_mesh(mesh_context, device_id, get_axis_index)
  if device_id_type is DeviceIdType.MESH:
    # Mesh means we are passed the mesh coordinates for the device
    device_ids = tree_util.tree_leaves(device_id)
    mesh_strides: tuple[int, ...]
    if mesh_context is None:
      mesh_strides = ()
    else:
      mesh_strides = mesh_context.mesh_strides
    if len(device_ids) != len(mesh_strides):
      raise ValueError(
          "Number of device ids must match the number of mesh axes, but got"
          f" {len(device_ids)} ids for a {len(mesh_strides)}D mesh."
      )

    if not device_ids:
      # If there are no device ids, then it is purely local communication.
      return None, non_mesh_axes
    return sum(a * b for a, b in zip(device_ids, mesh_strides)), non_mesh_axes
  elif device_id_type is DeviceIdType.LOGICAL:
    return device_id, non_mesh_axes
  raise NotImplementedError(f"Unsupported device id type: {device_id_type}")


delay_p = jax_core.Primitive("delay")
delay_p.multiple_results = True


class DelayEffect(effects.Effect):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide exactly one device id per mesh axis
2. Derive the ids from the mesh context (e.g., mesh_strides/get_axis_index) rather than hardcoding

Example fix

// before
# 3D mesh
device_ids=(i, j)
// after
device_ids=(i, j, k)

assert len(device_ids) == len(mesh_strides), (
    f"{len(device_ids)} ids for {len(mesh_strides)}D mesh")

When it happens

Trigger: Passing a tuple of device ids whose length differs from len(mesh_strides), e.g., 2 ids to a 3D mesh.

Common situations: Hardcoding device id tuples while the mesh configuration changes; passing partial mesh coordinates.

Related errors

• {axis} mixes JAX mesh and Pallas mesh grid axes
• {non_mesh_axes}
• `device_id_type` must be MESH if `device_id` is a dict, got:
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• {name} only accepts inputs that are unreduced. Got {aval.str

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