jax-ml/jax · error · ValueError

top_level_all_gather doesn't allow input {aval} to be unshar

Error message

top_level_all_gather doesn't allow input {aval} to be unsharded on dimension {axis} when {out_spec=}.

What it means

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.

Source

Thrown at jax/_src/shard_map.py:2229

  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')
    return x
  return api.jit(shard_map(f_shmap, out_specs=out_spec))(x)

def top_level_all_gather(xs, out_sharding, *, multi_dim: bool = False):
  if not get_abstract_mesh().are_all_axes_explicit:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change out_spec for that dimension to match the input (gather-only: prefix of in_spec or None)
2. Pre-shard the input on that dimension before calling (e.g. via jit in_shardings)

Example fix

# before
top_level_all_gather(x, NamedSharding(mesh, P('data', None)))  # x is P(None, None)

# after
top_level_all_gather(x, NamedSharding(mesh, P(None, None)))

in_p, out_p = list(x.sharding.spec), list(out_named.spec)
for ax,(i,o) in enumerate(safe_zip(in_p,out_p)):
    if i is None and o is not None:
        raise AssertionError(f'dim {ax}: cannot shard via all_gather (in={i}, out={o})')

When it happens

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

Common situations: Misreading the API as a general resharding utility instead of an all-gather.

Related errors

• out_sharding passed to top_level_all_gather cannot be {out_s
• Input sharding mesh {aval.sharding.mesh} should be equal to
• multiple dimensions cannot be all_gathered since multi_dim=F
• 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/5d2e70398b67b2ee. Report an issue: GitHub.