jax-ml/jax · error · ValueError

The size of all_to_all split_axis ({x.shape[split_axis]}) ha

Error message

The size of all_to_all split_axis ({x.shape[split_axis]}) has to be divisible by the size of the named axis {axis_name} ({group_size})

What it means

jax.lax.all_to_all with tiled=True splits x along split_axis across the mapped axis; the split dimension size must be divisible by the number of devices (group_size) so each device gets an equal tile.

Source

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

    the input ``x``.

    Otherwise array with shape similar to the input shape, except with split_axis
    divided by axis size and concat_axis multiplied by axis size.
  """
  return _all_to_all_is_async(x, axis_name, split_axis, concat_axis,
                              axis_index_groups=axis_index_groups, tiled=tiled,
                              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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the split axis up to the next multiple of group_size (and trim after)
2. Choose a split_axis whose size is divisible by the mapped-axis size
3. Adjust the mesh so the axis size divides the dimension evenly

Example fix

# before
y = jax.lax.all_to_all(x, 'i', split_axis=1, concat_axis=1, tiled=True)  # shape (m,6), 4 devices

# after
pad = (-x.shape[1]) % 4
y = jax.lax.all_to_all(jnp.pad(x, ((0,0),(0,pad))), 'i', 1, 1, tiled=True)[..., :x.shape[1]*2//2]  # trim as needed

assert x.shape[split_axis] % group_size == 0, (x.shape, group_size)

When it happens

Trigger: all_to_all(x, 'i', split_axis=1, concat_axis=1, tiled=True) where x.shape[1] (e.g. 6) is not divisible by the axis size (e.g. 4).

Common situations: Sequence/model dims not multiples of device count; changing mesh size without re-padding the feature dimension; hand-tiled SPMD code.

Related errors

• all_to_all requires the size of the mapped axis axis_name to
• {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/d648ad3e144920e2. Report an issue: GitHub.