jax-ml/jax · error · ValueError

multiple dimensions cannot be all_gathered since multi_dim=F

Error message

multiple dimensions cannot be all_gathered since multi_dim=False passed to `top_level_all_gather`. Got {in_spec=} and {out_spec=}

What it means

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.

Source

Thrown at jax/_src/shard_map.py:2224

  if aval.sharding.mesh != out_sh.mesh:
    raise ValueError(
        f'Input sharding mesh {aval.sharding.mesh} should be equal to'
        f' out_sharding mesh {out_sh.mesh}')

  in_spec = aval.sharding.spec
  out_spec = out_sh.spec._normalized_spec_for_aval(len(in_spec))
  if config.remove_size_one_mesh_axis_from_type.value:
    out_spec = remove_size_one_mesh_axis_from_spec(out_spec, out_sh.mesh)

  def f_shmap(x):
    # Maybe this can just be 1 AG where we gather in a new dim and then do
    # AG(new_dim) -> reshape -> transpose -> reshape but it might be expensive.
    count = 0
    for axis, (i, o) in enumerate(zip(in_spec.partitions, out_spec.partitions)):
      if i == o:
        continue
      if not multi_dim and count > 0:
        raise ValueError(
            "multiple dimensions cannot be all_gathered since multi_dim=False"
            f" passed to `top_level_all_gather`. Got {in_spec=} and {out_spec=}")
      count += 1
      if i is None:
        raise ValueError(
            f"top_level_all_gather doesn't allow input {aval} to be unsharded"
            f" on dimension {axis} when {out_spec=}.")
      i = i if isinstance(i, tuple) else (i,)
      o = o if o is None or isinstance(o, tuple) else (o,)
      if o is not None and i[:len(o)] != o:
        raise ValueError(
            'top_level_all_gather maintains `top_level_all_gather(x, ...) == x`'
            f" property. The {in_spec=} and {out_spec=} don't satisfy this"
            f' property. Please change your out_spec of array dimension {axis} so'
            " that it's a prefix of in_spec")
      axis_name = i if o is None else i[-len(o):]
      x = lax_parallel.all_gather(x, axis_name=axis_name, axis=axis,
                                  tiled=True, to='reduced')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass multi_dim=True to allow gathering across multiple dimensions
2. Or change out_sharding so only one dimension differs from the input's sharding

Example fix

# before
top_level_all_gather(x, NamedSharding(mesh, P('a','b')))

# after
top_level_all_gather(x, NamedSharding(mesh, P('a','b')), multi_dim=True)

in_p, out_p = x.sharding.spec, out_named.spec
diffs = sum(1 for i, o in zip(in_p, out_p) if i != o)
assert diffs <= 1 or multi_dim, f'{diffs} dims change; pass multi_dim=True'

When it happens

Trigger: 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.

Common situations: Trying to fully replicate an array sharded on multiple axes and forgetting the multi_dim flag.

Related errors

• out_sharding passed to top_level_all_gather cannot be {out_s
• Input sharding mesh {aval.sharding.mesh} should be equal to
• top_level_all_gather doesn't allow input {aval} to be unshar
• top_level_all_gather maintains `top_level_all_gather(x, ...)
• Mesh must be provided for shard_map with checkify.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f585b29cc0288227. Report an issue: GitHub.