jax-ml/jax · error · NotImplementedError

Failed to find assignment for logical_axis_index {logical_ax

Error message

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.

What it means

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).

Source

Thrown at jax/_src/mesh_utils.py:372

        if np.prod(c_axes) == logical_axis_size:
          assignment[logical_axis_index] = c_indices
          # Zero the assigned physical axes.
          assignable_physical_mesh = [
              0 if i in c_indices else v
              for i, v in enumerate(assignable_physical_mesh)
          ]
          break
      if assignment[logical_axis_index]:
        # We already found an assignment from one candidate above.
        break
    else:
      # If the num_axes for loop did not break, i.e. none of the candidates work
      # goto here with this while-else construct.
      if logical_axis_size > 1:
        if not allow_split_physical_axes:
          # Although this is now implemented, there are downstream tasks
          # counting on this being a NotImplementedError.
          raise NotImplementedError(
              'Failed to find assignment for logical_axis_index'
              f' {logical_axis_index} of size {logical_axis_size} with'
              f' 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.'
          )
        else:
          # We will try finding an assignment, even if that means splitting the
          # physical axes, which requires a more sophisticated implementation.
          return _create_device_mesh_for_nd_torus_splitting_axes(
              physical_mesh, mesh_shape
          )

  # Flatten the assignment, e.g., [(), (2,), (0, 1)] -> (2, 0, 1).
  transpose: list[int] = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set allow_split_physical_axes=True if splitting a physical axis across logical axes is acceptable
2. Redesign the logical mesh so each axis size equals a product of whole physical axis sizes (16=4x4 on 4x4x4)
3. Use create_hybrid_device_mesh for hosts+isf topologies which handles more layouts
4. Verify your device count: logical mesh total size must equal physical device count

Example fix

# before
mesh = create_device_mesh((4, 16))  # on 4x4x4 TPU -> NotImplementedError

# after
mesh = create_device_mesh((4, 16), allow_split_physical_axes=True)
# or: mesh = create_device_mesh((16, 4))  # 16=4x4, 4=4 works without splitting

import math, jax
dev_count = jax.device_count()
assert math.prod(logical_mesh_shape) == dev_count
# without splitting, each axis size must be a product of physical axis sizes:
def factorable(size, phys):
    from itertools import combinations
    for r in range(len(phys)+1):
        for c in combinations(phys, r):
            p = 1
            for x in c: p *= x
            if p == size: return True
    return False

try:
    mesh = create_device_mesh(logical_shape)
except NotImplementedError:
    mesh = create_device_mesh(logical_shape, allow_split_physical_axes=True)

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• QDWH implementation is only supported on TPU
• When passing host local inputs to pjit, devices connected to
• AbstractMesh does not implement {name}
• Memory space {self.memory_space} is not supported by mesh {s
• When `check_vma=True` on `jax.shard_map`, `manual_axis_type`

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e93d1c19e529b74a. Report an issue: GitHub.