{"record":{"id":"be90d6490828204d","repo":"jax-ml/jax","slug":"top-level-all-gather-maintains-top-level-all-gath","errorCode":null,"errorMessage":"top_level_all_gather maintains `top_level_all_gather(x, ...) == x` property. The {in_spec=} and {out_spec=} don't satisfy this property. Please change your out_spec of array dimension {axis} so that it's a prefix of in_spec","messagePattern":"top_level_all_gather maintains `top_level_all_gather\\(x, \\.\\.\\.\\) == x` property\\. The (.+?) and (.+?) don't satisfy this property\\. Please change your out_spec of array dimension (.+?) so that it's a prefix of in_spec","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2235,"sourceCode":"    # Maybe this can just be 1 AG where we gather in a new dim and then do\n    # AG(new_dim) -> reshape -> transpose -> reshape but it might be expensive.\n    count = 0\n    for axis, (i, o) in enumerate(zip(in_spec.partitions, out_spec.partitions)):\n      if i == o:\n        continue\n      if not multi_dim and count > 0:\n        raise ValueError(\n            \"multiple dimensions cannot be all_gathered since multi_dim=False\"\n            f\" passed to `top_level_all_gather`. Got {in_spec=} and {out_spec=}\")\n      count += 1\n      if i is None:\n        raise ValueError(\n            f\"top_level_all_gather doesn't allow input {aval} to be unsharded\"\n            f\" on dimension {axis} when {out_spec=}.\")\n      i = i if isinstance(i, tuple) else (i,)\n      o = o if o is None or isinstance(o, tuple) else (o,)\n      if o is not None and i[:len(o)] != o:\n        raise ValueError(\n            'top_level_all_gather maintains `top_level_all_gather(x, ...) == x`'\n            f\" property. The {in_spec=} and {out_spec=} don't satisfy this\"\n            f' property. Please change your out_spec of array dimension {axis} so'\n            \" that it's a prefix of in_spec\")\n      axis_name = i if o is None else i[-len(o):]\n      x = lax_parallel.all_gather(x, axis_name=axis_name, axis=axis,\n                                  tiled=True, to='reduced')\n    return x\n  return api.jit(shard_map(f_shmap, out_specs=out_spec))(x)\n\ndef top_level_all_gather(xs, out_sharding, *, multi_dim: bool = False):\n  if not get_abstract_mesh().are_all_axes_explicit:\n    raise ValueError(\n        'top_level_all_gather works when all mesh axes of context mesh are'\n        f' explicit. Got {get_abstract_mesh()}')\n  x_flat, treedef = tree_flatten(xs)\n  out_sharding_flat = api_util.flatten_axis_resources(\n      \"top_level_all_gather out_sharding\", treedef, out_sharding,","sourceCodeStart":2217,"sourceCodeEnd":2253,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2217-L2253","documentation":"top_level_all_gather guarantees top_level_all_gather(x, ...) == x (semantically the same global array, differently laid out). If the out_spec for a dim is not a prefix of the in_spec (or None), the memory layout change would alter semantics, so it's rejected.","triggerScenarios":"Input P(('data','model'),) with out_spec P('model',) — not a prefix — raising the error; correct out_spec is P(None,) or P(('data','model'),) subsets that are prefixes.","commonSituations":"Trying to reorder gathered sub-axes or drop a non-leading sub-axis of a tuple-sharded dimension.","solutions":["Make each dimension's out_spec either None (fully gathered) or a prefix of the in_spec sub-axis tuple","If you need a different axis order, do an explicit transpose/reshape after the gather"],"exampleFix":"# before\ntop_level_all_gather(x, NamedSharding(mesh, P('model',)))  # in: P(('data','model'),)\n\n# after\ny = top_level_all_gather(x, NamedSharding(mesh, P(None,)))\ny = transpose_for_model(y)  # reorder explicitly afterwards","handlingStrategy":"validation","validationCode":"for ax,(i,o) in enumerate(zip(in_spec, out_spec)):\n    it = i if isinstance(i, tuple) else ((i,) if i else ())\n    ot = o if isinstance(o, tuple) else ((o,) if o else ())\n    assert o is None or it[:len(ot)] == ot, f'dim {ax}: out_spec not a prefix of in_spec'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Only fully gather dims (out None) or trim trailing sub-axes; reorder explicitly after the gather"],"tags":["jax","all-gather","shard-map","partition-spec"],"backgroundTag":"jax-all-gather-reshaping-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}