{"record":{"id":"f585b29cc0288227","repo":"jax-ml/jax","slug":"multiple-dimensions-cannot-be-all-gathered-since-m","errorCode":null,"errorMessage":"multiple dimensions cannot be all_gathered since multi_dim=False passed to `top_level_all_gather`. Got {in_spec=} and {out_spec=}","messagePattern":"multiple dimensions cannot be all_gathered since multi_dim=False passed to `top_level_all_gather`\\. Got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2224,"sourceCode":"  if aval.sharding.mesh != out_sh.mesh:\n    raise ValueError(\n        f'Input sharding mesh {aval.sharding.mesh} should be equal to'\n        f' out_sharding mesh {out_sh.mesh}')\n\n  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')","sourceCodeStart":2206,"sourceCodeEnd":2242,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2206-L2242","documentation":"With multi_dim=False (the default), top_level_all_gather may gather along only one array dimension. If comparing in_spec vs out_spec shows two or more dimensions change sharding, this error is raised.","triggerScenarios":"Calling top_level_all_gather(x, NamedSharding(mesh, P('a','b'))) when x is sharded P(None,'b') — two dims differ — without multi_dim=True.","commonSituations":"Trying to fully replicate an array sharded on multiple axes and forgetting the multi_dim flag.","solutions":["Pass multi_dim=True to allow gathering across multiple dimensions","Or change out_sharding so only one dimension differs from the input's sharding"],"exampleFix":"# before\ntop_level_all_gather(x, NamedSharding(mesh, P('a','b')))\n\n# after\ntop_level_all_gather(x, NamedSharding(mesh, P('a','b')), multi_dim=True)","handlingStrategy":"validation","validationCode":"in_p, out_p = x.sharding.spec, out_named.spec\ndiffs = sum(1 for i, o in zip(in_p, out_p) if i != o)\nassert diffs <= 1 or multi_dim, f'{diffs} dims change; pass multi_dim=True'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default to multi_dim=True when fully replicating multi-axis-sharded arrays","Compute spec diffs in a helper before calling"],"tags":["jax","all-gather","shard-map","multi-dim"],"backgroundTag":"jax-multi-dim-collective-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}