{"record":{"id":"5dba05d1a517b697","repo":"jax-ml/jax","slug":"number-of-device-ids-must-match-the-number-of-mesh","errorCode":null,"errorMessage":"Number of device ids must match the number of mesh axes, but got {len(device_ids)} ids for a {len(mesh_strides)}D mesh.","messagePattern":"Number of device ids must match the number of mesh axes, but got (.+?) ids for a (.+?)D mesh\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":1282,"sourceCode":"  \"\"\"\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\n    if not device_ids:\n      # If there are no device ids, then it is purely local communication.\n      return None, non_mesh_axes\n    return sum(a * b for a, b in zip(device_ids, mesh_strides)), non_mesh_axes\n  elif device_id_type is DeviceIdType.LOGICAL:\n    return device_id, non_mesh_axes\n  raise NotImplementedError(f\"Unsupported device id type: {device_id_type}\")\n\n\ndelay_p = jax_core.Primitive(\"delay\")\ndelay_p.multiple_results = True\n\n\nclass DelayEffect(effects.Effect):","sourceCodeStart":1264,"sourceCodeEnd":1300,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L1264-L1300","documentation":"When device_id_type is MESH, the number of device ids must equal the mesh dimensionality; got fewer/more ids than mesh axes.","triggerScenarios":"Passing a tuple of device ids whose length differs from len(mesh_strides), e.g., 2 ids to a 3D mesh.","commonSituations":"Hardcoding device id tuples while the mesh configuration changes; passing partial mesh coordinates.","solutions":["Provide exactly one device id per mesh axis","Derive the ids from the mesh context (e.g., mesh_strides/get_axis_index) rather than hardcoding"],"exampleFix":"// before\n# 3D mesh\ndevice_ids=(i, j)\n// after\ndevice_ids=(i, j, k)","handlingStrategy":"validation","validationCode":"assert len(device_ids) == len(mesh_strides), (\n    f\"{len(device_ids)} ids for {len(mesh_strides)}D mesh\")","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive device ids from the mesh context instead of hardcoding tuples","Assert id count matches mesh rank in helpers"],"tags":["pallas","device-id","mesh","validation","jax"],"backgroundTag":"argument-count-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}