{"record":{"id":"6e08eeb2ee3aa44e","repo":"jax-ml/jax","slug":"all-to-all-requires-the-size-of-the-mapped-axis-ax","errorCode":null,"errorMessage":"all_to_all requires the size of the mapped axis axis_name to equal x.shape[split_axis], but they are {} and {} respectively.","messagePattern":"all_to_all requires the size of the mapped axis axis_name to equal x\\.shape\\[split_axis\\], but they are (.+?) and (.+?) respectively\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":610,"sourceCode":"                              is_async=False)\n\ndef _all_to_all_is_async(x, axis_name, split_axis, concat_axis, *,\n                         axis_index_groups=None, tiled=False, is_async=False):\n  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)\n  def bind(x, split_axis=split_axis, concat_axis=concat_axis):\n    split_axis = canonicalize_axis(split_axis, np.ndim(x))\n    concat_axis = canonicalize_axis(concat_axis, np.ndim(x))\n    group_size = _axis_size(axis_name, axis_index_groups)\n    if tiled:\n      if x.shape[split_axis] % group_size != 0:\n        raise ValueError(f\"The size of all_to_all split_axis ({x.shape[split_axis]}) \"\n                         f\"has to be divisible by the size of the named axis \"\n                         f\"{axis_name} ({group_size})\")\n    else:\n      if group_size != x.shape[split_axis]:\n        msg = (\"all_to_all requires the size of the mapped axis axis_name to \"\n               \"equal x.shape[split_axis], but they are {} and {} respectively.\")\n        raise ValueError(msg.format(group_size, x.shape[split_axis]))\n      if split_axis < concat_axis:\n        concat_axis += 1  # concat_axis gives a position _after_ split_axis is removed\n        x = lax.expand_dims(x, (concat_axis,))  # insert the new axis\n      elif split_axis == concat_axis:\n        pass\n      else:  # concat_axis < split_axis\n        x = lax.expand_dims(x, (concat_axis,))  # insert the new axis\n        split_axis += 1   # we have a new axis before split_axis now\n    x = insert_collective_pvary(axis_name, x)\n    prim = all_to_all_start_p if is_async else all_to_all_p\n    result = prim.bind(x, split_axis=split_axis, concat_axis=concat_axis,\n                               axis_name=axis_name,\n                               axis_index_groups=axis_index_groups,\n                               tiled=tiled)\n    if not tiled and split_axis != concat_axis:\n      result = lax.squeeze(result, (split_axis,))\n    return result\n","sourceCodeStart":592,"sourceCodeEnd":628,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L592-L628","documentation":"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.","triggerScenarios":"all_to_all(x, 'i', split_axis=0) where x.shape[0] is 3 but the 'i' axis has 8 devices (or vice versa).","commonSituations":"Assuming the split dim equals device count after reshapes; mismatch between the pmapped axis size and a hand-padded dimension; mesh reconfiguration.","solutions":["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)","Use tiled=True if the split dim is a multiple of, rather than equal to, the axis size","Recheck mesh definitions so the named axis size matches your data layout"],"exampleFix":"# before\ny = jax.lax.all_to_all(x, 'i', split_axis=1, concat_axis=1)  # x.shape[1]=16, axis size 8\n\n# after\nx2 = x.reshape(x.shape[0], 8, x.shape[1]//8)\ny = jax.lax.all_to_all(x2, 'i', split_axis=1, concat_axis=0).reshape(x.shape[0], -1)","handlingStrategy":"validation","validationCode":"group_size = jax.lax.psum(1, axis_name)\nassert group_size == x.shape[split_axis], (group_size, x.shape[split_axis])","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep one axis of size == device count for tiled=False all_to_all","Prefer tiled=True when dims are only multiples"],"tags":["jax","all-to-all","shape","spmd"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}