jax-ml/jax · error · ValueError

`device_id_type` must be MESH if `device_id` is a dict, got:

Error message

`device_id_type` must be MESH if `device_id` is a dict, got: {device_id_type = }.

What it means

device_id_to_logical requires DeviceIdType.MESH when device_id is passed as a dict (mesh-coordinate form). Any other device_id_type with a dict device_id is a ValueError.

Source

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

  return tuple(device_id), non_mesh_axes


def device_id_to_logical(
    mesh_context: pallas_utils.MeshInfo | None,
    device_id: Any,
    device_id_type: DeviceIdType,
    get_axis_index: Callable[[Any], Any],
) -> tuple[Any | None, dict[Any, Any]]:
  """Normalizes a device id into a logical device id and axes that don't correspond to JAX mesh axes.

  The indexing implied by the returned axis dict should be handled by the
  caller. If there are no cross-device operations, then the returned logical
  device id will be None.
  """
  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."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set device_id_type=DeviceIdType.MESH when passing a dict device_id
Or pass a plain integer device id with LOGICAL type

Example fix

// before
signal_remote(sem, device_id={'data': i})
// after
signal_remote(sem, device_id={'data': i}, device_id_type=DeviceIdType.MESH)

Defensive patterns

Strategy: type-guard

Validate before calling

if isinstance(device_id, dict):
    assert device_id_type is DeviceIdType.MESH, "dict device_id requires DeviceIdType.MESH"

Type guard

def check_device_id(device_id, device_id_type) -> bool:
    if isinstance(device_id, dict):
        return device_id_type is DeviceIdType.MESH
    return device_id_type in (DeviceIdType.MESH, DeviceIdType.LOGICAL)

Prevention

Pair dict device ids with DeviceIdType.MESH always
Validate device_id/device_id_type combos in a wrapper function

When it happens

Trigger: Calling an API that routes through device_id_to_logical with device_id={'x': 0} but device_id_type=DeviceIdType.LOGICAL (or default).

Common situations: Passing mesh coordinate dicts from a Mesh context while leaving device_id_type at LOGICAL; API wrappers that always pass dicts.

Related errors

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