{"record":{"id":"5d2e70398b67b2ee","repo":"jax-ml/jax","slug":"top-level-all-gather-doesn-t-allow-input-aval-to","errorCode":null,"errorMessage":"top_level_all_gather doesn't allow input {aval} to be unsharded on dimension {axis} when {out_spec=}.","messagePattern":"top_level_all_gather doesn't allow input (.+?) to be unsharded on dimension (.+?) when (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2229,"sourceCode":"  in_spec = aval.sharding.spec\n  out_spec = out_sh.spec._normalized_spec_for_aval(len(in_spec))\n  if config.remove_size_one_mesh_axis_from_type.value:\n    out_spec = remove_size_one_mesh_axis_from_spec(out_spec, out_sh.mesh)\n\n  def f_shmap(x):\n    # 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:","sourceCodeStart":2211,"sourceCodeEnd":2247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2211-L2247","documentation":"top_level_all_gather can only gather, not shard: if the input is unsharded (in_spec None) on a dimension where out_spec expects sharding, gathering cannot produce it and the error is raised.","triggerScenarios":"Input replicated on dim i (P(None,...)) but out_sharding requests P('data') on that dim — i.e. asking the gather to also scatter/shard.","commonSituations":"Misreading the API as a general resharding utility instead of an all-gather.","solutions":["Change out_spec for that dimension to match the input (gather-only: prefix of in_spec or None)","Pre-shard the input on that dimension before calling (e.g. via jit in_shardings)"],"exampleFix":"# before\ntop_level_all_gather(x, NamedSharding(mesh, P('data', None)))  # x is P(None, None)\n\n# after\ntop_level_all_gather(x, NamedSharding(mesh, P(None, None)))","handlingStrategy":"validation","validationCode":"in_p, out_p = list(x.sharding.spec), list(out_named.spec)\nfor ax,(i,o) in enumerate(safe_zip(in_p,out_p)):\n    if i is None and o is not None:\n        raise AssertionError(f'dim {ax}: cannot shard via all_gather (in={i}, out={o})')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Treat top_level_all_gather as gather-only: out specs may only remove sharding, never add it"],"tags":["jax","all-gather","shard-map","resharding"],"backgroundTag":"jax-all-gather-reshaping-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}