{"record":{"id":"d648ad3e144920e2","repo":"jax-ml/jax","slug":"the-size-of-all-to-all-split-axis-x-shape-split","errorCode":null,"errorMessage":"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})","messagePattern":"The size of all_to_all split_axis \\((.+?)\\) has to be divisible by the size of the named axis (.+?) \\((.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":603,"sourceCode":"    the input ``x``.\n\n    Otherwise array with shape similar to the input shape, except with split_axis\n    divided by axis size and concat_axis multiplied by axis size.\n  \"\"\"\n  return _all_to_all_is_async(x, axis_name, split_axis, concat_axis,\n                              axis_index_groups=axis_index_groups, tiled=tiled,\n                              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,","sourceCodeStart":585,"sourceCodeEnd":621,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L585-L621","documentation":"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.","triggerScenarios":"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).","commonSituations":"Sequence/model dims not multiples of device count; changing mesh size without re-padding the feature dimension; hand-tiled SPMD code.","solutions":["Pad the split axis up to the next multiple of group_size (and trim after)","Choose a split_axis whose size is divisible by the mapped-axis size","Adjust the mesh so the axis size divides the dimension evenly"],"exampleFix":"# before\ny = jax.lax.all_to_all(x, 'i', split_axis=1, concat_axis=1, tiled=True)  # shape (m,6), 4 devices\n\n# after\npad = (-x.shape[1]) % 4\ny = 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","handlingStrategy":"validation","validationCode":"assert x.shape[split_axis] % group_size == 0, (x.shape, group_size)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Choose split dims that are multiples of device count","Pad dims defensively before all_to_all"],"tags":["jax","all-to-all","shape","spmd"],"backgroundTag":"dimension-not-divisible","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}