jax-ml/jax · error · ValueError

all_to_all requires the size of the mapped axis axis_name to

Error message

all_to_all requires the size of the mapped axis axis_name to equal x.shape[split_axis], but they are {} and {} respectively.

What it means

jax.lax.all_to_all with tiled=False requires an exact match: the mapped axis size must equal x.shape[split_axis], because each device contributes exactly one slice of the split dimension.

Source

Thrown at jax/_src/lax/parallel.py:610

                              is_async=False)

def _all_to_all_is_async(x, axis_name, split_axis, concat_axis, *,
                         axis_index_groups=None, tiled=False, is_async=False):
  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)
  def bind(x, split_axis=split_axis, concat_axis=concat_axis):
    split_axis = canonicalize_axis(split_axis, np.ndim(x))
    concat_axis = canonicalize_axis(concat_axis, np.ndim(x))
    group_size = _axis_size(axis_name, axis_index_groups)
    if tiled:
      if x.shape[split_axis] % group_size != 0:
        raise ValueError(f"The size of all_to_all split_axis ({x.shape[split_axis]}) "
                         f"has to be divisible by the size of the named axis "
                         f"{axis_name} ({group_size})")
    else:
      if group_size != x.shape[split_axis]:
        msg = ("all_to_all requires the size of the mapped axis axis_name to "
               "equal x.shape[split_axis], but they are {} and {} respectively.")
        raise ValueError(msg.format(group_size, x.shape[split_axis]))
      if split_axis < concat_axis:
        concat_axis += 1  # concat_axis gives a position _after_ split_axis is removed
        x = lax.expand_dims(x, (concat_axis,))  # insert the new axis
      elif split_axis == concat_axis:
        pass
      else:  # concat_axis < split_axis
        x = lax.expand_dims(x, (concat_axis,))  # insert the new axis
        split_axis += 1   # we have a new axis before split_axis now
    x = insert_collective_pvary(axis_name, x)
    prim = all_to_all_start_p if is_async else all_to_all_p
    result = prim.bind(x, split_axis=split_axis, concat_axis=concat_axis,
                               axis_name=axis_name,
                               axis_index_groups=axis_index_groups,
                               tiled=tiled)
    if not tiled and split_axis != concat_axis:
      result = lax.squeeze(result, (split_axis,))
    return result

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/transpose so the split dimension equals the axis size (e.g. add a leading axis of size axis_size and split a different dim)
2. Use tiled=True if the split dim is a multiple of, rather than equal to, the axis size
3. Recheck mesh definitions so the named axis size matches your data layout

Example fix

# before
y = jax.lax.all_to_all(x, 'i', split_axis=1, concat_axis=1)  # x.shape[1]=16, axis size 8

# after
x2 = x.reshape(x.shape[0], 8, x.shape[1]//8)
y = jax.lax.all_to_all(x2, 'i', split_axis=1, concat_axis=0).reshape(x.shape[0], -1)

group_size = jax.lax.psum(1, axis_name)
assert group_size == x.shape[split_axis], (group_size, x.shape[split_axis])

When it happens

Trigger: all_to_all(x, 'i', split_axis=0) where x.shape[0] is 3 but the 'i' axis has 8 devices (or vice versa).

Common situations: Assuming the split dim equals device count after reshapes; mismatch between the pmapped axis size and a hand-padded dimension; mesh reconfiguration.

Related errors

• The size of all_to_all split_axis ({x.shape[split_axis]}) ha
• {name} was requested to map its argument along axis {axis},
• {''.join(msg)[:-2]}
• jvp called with different primal and tangent shapes;Got prim
• callbacks are only supported in spmd computations when all m

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