jax-ml/jax · error · ValueError

Replica groups must be equally sized

Error message

Replica groups must be equally sized

What it means

all_to_all splits each replica's tensor across the group and recombines; XLA's AllToAll requires every replica group to have identical size. The lowering checks all groups equal the first group's size before emitting the op.

Source

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

  new_shape = list(x.shape)
  new_shape[axis:axis+2] = [x.shape[axis] * x.shape[axis + 1]]
  return x.reshape(new_shape)

def _all_to_all_lowering(
    ctx, x, *, split_axis, concat_axis, axis_name, axis_index_groups, tiled,
    is_async=False
):
  del tiled  # expand_dims and squeeze is done in `all_to_all` if `True`
  # Workaround for AllToAll not being implemented on CPU.
  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name,
                                   axis_index_groups)
  if not is_async and len(replica_groups[0]) == 1:
    # TODO(mwhittaker): This optimization doesn't play well with async
    # collectives. Support it; or optimize it in XLA.
    return [x]
  split_count = len(replica_groups[0])
  if not all(split_count == len(g) for g in replica_groups):
    raise ValueError('Replica groups must be equally sized')
  is_spmd = isinstance(
      ctx.module_context.axis_context,
      (SPMDAxisContext, ShardingContext),
  )
  if is_spmd:
    # We want to emit the all-gather with global device IDs and a
    # channel ID, as otherwise it interprets the devices as replicas instead
    # of partitions - and XLA is configured with only a single replica.
    channel_handle = hlo.ChannelHandle.get(mlir.COLLECTIVE_CHANNEL_ID,
                                           mlir.DEVICE_TO_DEVICE_TYPE)
    other_args: dict[str, Any] = dict(channel_handle=channel_handle)
  else:
    other_args = {}

  replica_groups_attr = _try_mesh_axes_replica_group(
      ctx, axis_name, axis_index_groups
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all axis_index_groups the same size (pad or regroup)
2. Use mesh axis sizes that divide group sizes evenly
3. Drop axis_index_groups for full-mesh all_to_all

Example fix

// before
lax.all_to_all(x, 'i', 0, 0, axis_index_groups=[[0,1],[2]])
// after
lax.all_to_all(x, 'i', 0, 0, axis_index_groups=[[0,1],[2,3]])

def check_groups(groups):
    assert groups is None or len({len(g) for g in groups}) == 1, 'replica groups must be equally sized'

When it happens

Trigger: Calling jax.lax.all_to_all with axis_index_groups of differing lengths, or a mesh grouping that produces unequal replica groups.

Common situations: Custom axis_index_groups for sub-group all-to-all; uneven device partitions on heterogeneous clusters.

Related errors

• {op_name} sources and destinations must be unique, got {}.
• {name} only accepts inputs that are unreduced. Got {aval.str
• {full_name} must be a pytree prefix with bool leaves or a tu
• unbound axis name: {axis_name}
• Axes mentioned in `manual_axis_type` field of ShapedArray sh

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