{"record":{"id":"b6d333e1b0c6023c","repo":"jax-ml/jax","slug":"device-id-type-must-be-mesh-if-device-id-is-a-b6d333","errorCode":null,"errorMessage":"`device_id_type` must be MESH if `device_id` is a dict, got: {device_id_type = }.","messagePattern":"`device_id_type` must be MESH if `device_id` is a dict, got: (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":1268,"sourceCode":"  return tuple(device_id), non_mesh_axes\n\n\ndef device_id_to_logical(\n    mesh_context: pallas_utils.MeshInfo | None,\n    device_id: Any,\n    device_id_type: DeviceIdType,\n    get_axis_index: Callable[[Any], Any],\n) -> tuple[Any | None, dict[Any, Any]]:\n  \"\"\"Normalizes a device id into a logical device id and axes that don't correspond to JAX mesh axes.\n\n  The indexing implied by the returned axis dict should be handled by the\n  caller. If there are no cross-device operations, then the returned logical\n  device id will be None.\n  \"\"\"\n  non_mesh_axes = {}\n  if isinstance(device_id, dict):\n    if device_id_type is not DeviceIdType.MESH:\n      raise ValueError(\n          \"`device_id_type` must be MESH if `device_id` is a dict,\"\n          f\" got: {device_id_type = }.\"\n      )\n    device_id, non_mesh_axes = _device_id_dict_to_mesh(mesh_context, device_id, get_axis_index)\n  if device_id_type is DeviceIdType.MESH:\n    # Mesh means we are passed the mesh coordinates for the device\n    device_ids = tree_util.tree_leaves(device_id)\n    mesh_strides: tuple[int, ...]\n    if mesh_context is None:\n      mesh_strides = ()\n    else:\n      mesh_strides = mesh_context.mesh_strides\n    if len(device_ids) != len(mesh_strides):\n      raise ValueError(\n          \"Number of device ids must match the number of mesh axes, but got\"\n          f\" {len(device_ids)} ids for a {len(mesh_strides)}D mesh.\"\n      )\n","sourceCodeStart":1250,"sourceCodeEnd":1286,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L1250-L1286","documentation":"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.","triggerScenarios":"Calling an API that routes through device_id_to_logical with device_id={'x': 0} but device_id_type=DeviceIdType.LOGICAL (or default).","commonSituations":"Passing mesh coordinate dicts from a Mesh context while leaving device_id_type at LOGICAL; API wrappers that always pass dicts.","solutions":["Set device_id_type=DeviceIdType.MESH when passing a dict device_id","Or pass a plain integer device id with LOGICAL type"],"exampleFix":"// before\nsignal_remote(sem, device_id={'data': i})\n// after\nsignal_remote(sem, device_id={'data': i}, device_id_type=DeviceIdType.MESH)","handlingStrategy":"type-guard","validationCode":"if isinstance(device_id, dict):\n    assert device_id_type is DeviceIdType.MESH, \"dict device_id requires DeviceIdType.MESH\"","typeGuard":"def check_device_id(device_id, device_id_type) -> bool:\n    if isinstance(device_id, dict):\n        return device_id_type is DeviceIdType.MESH\n    return device_id_type in (DeviceIdType.MESH, DeviceIdType.LOGICAL)","tryCatchPattern":null,"preventionTips":["Pair dict device ids with DeviceIdType.MESH always","Validate device_id/device_id_type combos in a wrapper function"],"tags":["pallas","device-id","validation","jax"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}