{"record":{"id":"e93d1c19e529b74a","repo":"jax-ml/jax","slug":"failed-to-find-assignment-for-logical-axis-index","errorCode":null,"errorMessage":"Failed to find assignment for logical_axis_index {logical_axis_index} of size {logical_axis_size} with remaining assignable mesh {assignable_physical_mesh}. The size of each axis in your logical mesh must be equal to the product of some subset of the physical mesh axis sizes. E.g. logical mesh (4, 16) is compatible with physical mesh 4x4x4 since 4=4 and 16=4x4. If you want to split physical axes, set  allow_split_physical_axes to True.","messagePattern":"Failed to find assignment for logical_axis_index (.+?) of size (.+?) with remaining assignable mesh (.+?)\\. The size of each axis in your logical mesh must be equal to the product of some subset of the physical mesh axis sizes\\. E\\.g\\. logical mesh \\(4, 16\\) is compatible with physical mesh 4x4x4 since 4=4 and 16=4x4\\. If you want to split physical axes, set  allow_split_physical_axes to True\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh_utils.py","lineNumber":372,"sourceCode":"        if np.prod(c_axes) == logical_axis_size:\n          assignment[logical_axis_index] = c_indices\n          # Zero the assigned physical axes.\n          assignable_physical_mesh = [\n              0 if i in c_indices else v\n              for i, v in enumerate(assignable_physical_mesh)\n          ]\n          break\n      if assignment[logical_axis_index]:\n        # We already found an assignment from one candidate above.\n        break\n    else:\n      # If the num_axes for loop did not break, i.e. none of the candidates work\n      # goto here with this while-else construct.\n      if logical_axis_size > 1:\n        if not allow_split_physical_axes:\n          # Although this is now implemented, there are downstream tasks\n          # counting on this being a NotImplementedError.\n          raise NotImplementedError(\n              'Failed to find assignment for logical_axis_index'\n              f' {logical_axis_index} of size {logical_axis_size} with'\n              f' remaining assignable mesh {assignable_physical_mesh}. The size'\n              ' of each axis in your logical mesh must be equal to the product'\n              ' of some subset of the physical mesh axis sizes. E.g. logical'\n              ' mesh (4, 16) is compatible with physical mesh 4x4x4 since 4=4'\n              ' and 16=4x4. If you want to split physical axes, set '\n              ' allow_split_physical_axes to True.'\n          )\n        else:\n          # We will try finding an assignment, even if that means splitting the\n          # physical axes, which requires a more sophisticated implementation.\n          return _create_device_mesh_for_nd_torus_splitting_axes(\n              physical_mesh, mesh_shape\n          )\n\n  # Flatten the assignment, e.g., [(), (2,), (0, 1)] -> (2, 0, 1).\n  transpose: list[int] = []","sourceCodeStart":354,"sourceCodeEnd":390,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh_utils.py#L354-L390","documentation":"When building an N-D torus device mesh, mesh_utils must assign each logical axis a product of remaining physical axis sizes. Without allow_split_physical_axes=True, it only tries subsets of whole physical axes; if logical_axis_size (e.g. 16) can't be formed as a product of the remaining assignable physical axes, it raises NotImplementedError (kept for backward compatibility).","triggerScenarios":"create_device_mesh(logical_mesh) where a logical axis size doesn't factor into whole physical axes — e.g. logical mesh (4,16) on physical 4x4x4 with default allow_split_physical_axes=False; sizes like 8 on a 4x4x4 torus (8=4x2 requires splitting a 4).","commonSituations":"TPU pod-slice topologies (4x4x4, 8x8x8) where user-requested logical meshes don't align with physical axis sizes; H100 clusters with NVLink groups of 8; the classic XLA error text carried over for logical meshes incompatible with the torus.","solutions":["Set allow_split_physical_axes=True if splitting a physical axis across logical axes is acceptable","Redesign the logical mesh so each axis size equals a product of whole physical axis sizes (16=4x4 on 4x4x4)","Use create_hybrid_device_mesh for hosts+isf topologies which handles more layouts","Verify your device count: logical mesh total size must equal physical device count"],"exampleFix":"# before\nmesh = create_device_mesh((4, 16))  # on 4x4x4 TPU -> NotImplementedError\n\n# after\nmesh = create_device_mesh((4, 16), allow_split_physical_axes=True)\n# or: mesh = create_device_mesh((16, 4))  # 16=4x4, 4=4 works without splitting","handlingStrategy":"fallback","validationCode":"import math, jax\ndev_count = jax.device_count()\nassert math.prod(logical_mesh_shape) == dev_count\n# without splitting, each axis size must be a product of physical axis sizes:\ndef factorable(size, phys):\n    from itertools import combinations\n    for r in range(len(phys)+1):\n        for c in combinations(phys, r):\n            p = 1\n            for x in c: p *= x\n            if p == size: return True\n    return False","typeGuard":null,"tryCatchPattern":"try:\n    mesh = create_device_mesh(logical_shape)\nexcept NotImplementedError:\n    mesh = create_device_mesh(logical_shape, allow_split_physical_axes=True)","preventionTips":["Design logical axes as products of physical torus axis sizes","Pass allow_split_physical_axes=True when topology is irregular","Validate total size equals device count before calling"],"tags":["jax","mesh","tpu","topology","not-implemented"],"backgroundTag":"logical-physical-mesh-incompatible","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}