{"record":{"id":"cb4d5e38c803ebe8","repo":"jax-ml/jax","slug":"abstractmesh-should-be-of-the-same-size-across-all","errorCode":null,"errorMessage":"AbstractMesh should be of the same size across all shardings. Got {abstract_mesh} and {sh.mesh}","messagePattern":"AbstractMesh should be of the same size across all shardings\\. Got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":567,"sourceCode":"\ndef _get_and_check_device_assignment(\n    shardings: Iterable[ShardingInfo],\n    ctx_mesh: Mesh | AbstractMesh,\n) -> tuple[xc.Client, tuple[xc.Device, ...] | None, int]:\n  first_sharding_info = None\n  abstract_mesh = (\n      ctx_mesh if not ctx_mesh.empty and isinstance(ctx_mesh, AbstractMesh)\n      else None)\n  any_concrete_sharding = (\n      True if not ctx_mesh.empty and isinstance(ctx_mesh, Mesh) else False)\n\n  for sh, s_type, source_info in shardings:\n    if isinstance(sh, UnspecifiedValue):\n      continue\n    elif isinstance(sh, NamedSharding) and isinstance(sh.mesh, AbstractMesh):\n      if (abstract_mesh is not None and not sh.mesh.empty and\n          abstract_mesh.size != sh.mesh.size):\n        raise ValueError(\"AbstractMesh should be of the same size across all \"\n                         f\"shardings. Got {abstract_mesh} and {sh.mesh}\")\n      abstract_mesh = sh.mesh\n    else:\n      any_concrete_sharding = True\n      arr_device_assignment = sh._device_assignment\n      if first_sharding_info is None:\n        first_sharding_info = (arr_device_assignment, s_type, source_info)\n      if ctx_mesh.empty:\n        if first_sharding_info[0] != arr_device_assignment:\n          raise stages.DeviceAssignmentMismatchError([\n              stages.DeviceAssignmentMismatch(*first_sharding_info),\n              stages.DeviceAssignmentMismatch(\n                  arr_device_assignment, s_type, source_info)])\n      elif isinstance(ctx_mesh, AbstractMesh):\n        if ctx_mesh.size != len(arr_device_assignment):\n          raise stages.DeviceAssignmentMismatchError([\n              stages.DeviceAssignmentMismatch(\n                  ctx_mesh.size, stages.MismatchType.CONTEXT_DEVICES, None),","sourceCodeStart":549,"sourceCodeEnd":585,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L549-L585","documentation":"When a computation mixes NamedShardings over an AbstractMesh, JAX requires all such meshes to have the same total size (number of devices). If shardings in the same lowering carry AbstractMeshes of differing sizes, device assignment is ambiguous and JAX raises during _get_and_check_device_assignment (used by lower_sharding_computation).","triggerScenarios":"Constructing a jitted function (pjit) where different arguments' NamedShardings use AbstractMesh objects with different total sizes — e.g. one mesh 2x4 and another 8x1, or one empty and one populated inconsistently; mixing out_shardings and in_shardings with different abstract meshes.","commonSituations":"Building shardings from multiple user-created Mesh/AbstractMesh objects; refactoring mesh shapes mid-experiment; passing an out_sharding from an old mesh with new in_shardings.","solutions":["Make all AbstractMeshes used for in/out shardings the same size (same total device count)","Derive all shardings from a single shared Mesh/AbstractMesh object","Check mesh.empty cases: ensure you aren't mixing an empty AbstractMesh placeholder with a sized one inconsistently"],"exampleFix":"# before\nmesh_a = Mesh(devices8, axis_names=('dp',))\nmesh_b = Mesh(devices4, axis_names=('dp',))\njit(f, in_shardings=NamedSharding(mesh_a, P('dp')), out_shardings=NamedSharding(mesh_b, P('dp')))\n\n# after\nmesh = Mesh(devices8, axis_names=('dp',))\njit(f, in_shardings=NamedSharding(mesh, P('dp')), out_shardings=NamedSharding(mesh, P('dp')))","handlingStrategy":"validation","validationCode":"meshes = [s.mesh for s in shardings if isinstance(s, NamedSharding) and isinstance(s.mesh, AbstractMesh)]\nsizes = {m.size for m in meshes if not m.empty}\nassert len(sizes) <= 1, f'AbstractMesh sizes differ: {sizes}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Construct all NamedShardings from one shared Mesh object","Assert mesh size consistency in setup for sharded entry points"],"tags":["jax","sharding","abstract-mesh","pjit","mesh-size-mismatch"],"backgroundTag":"sharding-mesh-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}