{"record":{"id":"4b94710e99745cc3","repo":"jax-ml/jax","slug":"abstractmesh-size-abstract-mesh-size-does-not-m","errorCode":null,"errorMessage":"AbstractMesh size: {abstract_mesh.size} does not match the device assignment size: {len(device_assignment)}","messagePattern":"AbstractMesh size: (.+?) does not match the device assignment size: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":610,"sourceCode":"                  ctx_mesh._flat_devices_tuple,\n                  stages.MismatchType.CONTEXT_DEVICES, None),\n              stages.DeviceAssignmentMismatch(\n                  arr_device_assignment, s_type, source_info)])\n\n  device_assignment: tuple[xc.Device, ...]\n  if (first_sharding_info is None and not ctx_mesh.empty and\n      isinstance(ctx_mesh, Mesh)):\n    device_assignment = ctx_mesh._flat_devices_tuple\n  elif first_sharding_info is None:\n    device_assignment = (get_default_device(),)\n  else:\n    device_assignment = first_sharding_info[0]  # pyrefly: ignore[bad-assignment]\n\n  backend = xb.get_device_backend(device_assignment[0])\n\n  if (any_concrete_sharding and abstract_mesh is not None and\n      len(device_assignment) != abstract_mesh.size):\n    raise ValueError(\n        f\"AbstractMesh size: {abstract_mesh.size} does not match the\"\n        f\" device assignment size: {len(device_assignment)}\")\n\n  if any_concrete_sharding or abstract_mesh is None:\n    return backend, device_assignment, len(device_assignment)\n  else:\n    return backend, None, abstract_mesh.size\n\nMaybeSharding = JSharding | UnspecifiedValue\n\n\ndef prune_unused_inputs(\n    jaxpr: core.Jaxpr,\n) -> tuple[core.Jaxpr, set[int], set[int]]:\n  used_outputs = [True] * len(jaxpr.outvars)\n  new_jaxpr, used_consts, used_inputs = pe.dce_jaxpr_consts(jaxpr, used_outputs)\n  kept_const_idx = {i for i, b in enumerate(used_consts) if b}\n  kept_var_idx = {i for i, b in enumerate(used_inputs) if b}","sourceCodeStart":592,"sourceCodeEnd":628,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L592-L628","documentation":"When a computation mixes concrete shardings with NamedShardings over an AbstractMesh, JAX takes the device assignment from the concrete sharding and checks it against the AbstractMesh's total size; they must match. A mismatch means the sharded computation would be built for a different device set than the mesh declares.","triggerScenarios":"pjit with some args having concrete device shardings (e.g. SingleDeviceSharding, GSPMDSharding with a device list of N devices) and others using NamedSharding over an AbstractMesh of size M != N.","commonSituations":"Mixing concrete device assignments (e.g. from jax.device_put results or explicit GSPMDSharding) with abstract-mesh-based NamedShardings on other arguments/outputs; partial migration to AbstractMesh-based APIs.","solutions":["Make the concrete device_assignment length equal the AbstractMesh size (same device set)","Prefer all-abstract or all-concrete shardings rather than mixing","If mixing intentionally, construct the AbstractMesh over exactly the devices in the concrete assignment"],"exampleFix":"# before\nmesh = AbstractMesh((8,), axis_names=('dp',))  # size 8\ndevs = jax.devices()[:4]\njit(f, in_shardings=(GSPMDSharding(devs, ...), NamedSharding(mesh, P('dp'))))\n\n# after\nmesh = AbstractMesh((4,), axis_names=('dp',))  # matches len(devs)\njit(f, in_shardings=(GSPMDSharding(devs, ...), NamedSharding(mesh, P('dp'))))","handlingStrategy":"validation","validationCode":"abstract_sizes = [s.mesh.size for s in shardings if isinstance(s, NamedSharding) and isinstance(s.mesh, AbstractMesh)]\nconcrete = [s._device_assignment for s in shardings if not isinstance(s, NamedSharding)]\nif abstract_sizes and concrete:\n    assert all(len(d) == abstract_sizes[0] for d in concrete)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't mix concrete device shardings with AbstractMesh shardings on one call","Build the AbstractMesh over exactly the devices used in any concrete assignment"],"tags":["jax","sharding","abstract-mesh","device-assignment","pjit"],"backgroundTag":"sharding-mesh-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}