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

Guard in device_id_to_logical: passing device_id as a dict is the mesh-coordinates form, which is only interpreted when device_id_type is DeviceIdType.MESH. Supplying a dict with any other device_id_type is rejected since the dict's meaning would be undefined.

Solutions

1. Set device_id_type=DeviceIdType.MESH when passing a dict device_id
2. 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)

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

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)

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

• Number of device ids must match the number of mesh axes…
• All elements in the grid must be strictly positive, but got
• mixes JAX mesh and Pallas mesh grid axes
• should not be greater than
• Can't instantiate with arguments.

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

Appendix: 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)