{"record":{"id":"76fbb14a57bdcf5b","repo":"jax-ml/jax","slug":"replica-groups-must-be-equally-sized","errorCode":null,"errorMessage":"Replica groups must be equally sized","messagePattern":"Replica groups must be equally sized","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1434,"sourceCode":"  new_shape = list(x.shape)\n  new_shape[axis:axis+2] = [x.shape[axis] * x.shape[axis + 1]]\n  return x.reshape(new_shape)\n\ndef _all_to_all_lowering(\n    ctx, x, *, split_axis, concat_axis, axis_name, axis_index_groups, tiled,\n    is_async=False\n):\n  del tiled  # expand_dims and squeeze is done in `all_to_all` if `True`\n  # Workaround for AllToAll not being implemented on CPU.\n  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name,\n                                   axis_index_groups)\n  if not is_async and len(replica_groups[0]) == 1:\n    # TODO(mwhittaker): This optimization doesn't play well with async\n    # collectives. Support it; or optimize it in XLA.\n    return [x]\n  split_count = len(replica_groups[0])\n  if not all(split_count == len(g) for g in replica_groups):\n    raise ValueError('Replica groups must be equally sized')\n  is_spmd = isinstance(\n      ctx.module_context.axis_context,\n      (SPMDAxisContext, ShardingContext),\n  )\n  if is_spmd:\n    # We want to emit the all-gather with global device IDs and a\n    # channel ID, as otherwise it interprets the devices as replicas instead\n    # of partitions - and XLA is configured with only a single replica.\n    channel_handle = hlo.ChannelHandle.get(mlir.COLLECTIVE_CHANNEL_ID,\n                                           mlir.DEVICE_TO_DEVICE_TYPE)\n    other_args: dict[str, Any] = dict(channel_handle=channel_handle)\n  else:\n    other_args = {}\n\n  replica_groups_attr = _try_mesh_axes_replica_group(\n      ctx, axis_name, axis_index_groups\n  )\n","sourceCodeStart":1416,"sourceCodeEnd":1452,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1416-L1452","documentation":"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.","triggerScenarios":"Calling jax.lax.all_to_all with axis_index_groups of differing lengths, or a mesh grouping that produces unequal replica groups.","commonSituations":"Custom axis_index_groups for sub-group all-to-all; uneven device partitions on heterogeneous clusters.","solutions":["Make all axis_index_groups the same size (pad or regroup)","Use mesh axis sizes that divide group sizes evenly","Drop axis_index_groups for full-mesh all_to_all"],"exampleFix":"// before\nlax.all_to_all(x, 'i', 0, 0, axis_index_groups=[[0,1],[2]])\n// after\nlax.all_to_all(x, 'i', 0, 0, axis_index_groups=[[0,1],[2,3]])","handlingStrategy":"validation","validationCode":"def check_groups(groups):\n    assert groups is None or len({len(g) for g in groups}) == 1, 'replica groups must be equally sized'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive groups from divisors of axis size"],"tags":["jax","all-to-all","collectives","validation"],"backgroundTag":"unequal-group-sizes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}