jax-ml/jax · error · ValueError

top_level_all_gather maintains `top_level_all_gather(x, ...)

Error message

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

What it means

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.

Source

Thrown at jax/_src/shard_map.py:2235

    # 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:
    raise ValueError(
        'top_level_all_gather works when all mesh axes of context mesh are'
        f' explicit. Got {get_abstract_mesh()}')
  x_flat, treedef = tree_flatten(xs)
  out_sharding_flat = api_util.flatten_axis_resources(
      "top_level_all_gather out_sharding", treedef, out_sharding,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make each dimension's out_spec either None (fully gathered) or a prefix of the in_spec sub-axis tuple
  2. If you need a different axis order, do an explicit transpose/reshape after the gather

Example fix

# before
top_level_all_gather(x, NamedSharding(mesh, P('model',)))  # in: P(('data','model'),)

# after
y = top_level_all_gather(x, NamedSharding(mesh, P(None,)))
y = transpose_for_model(y)  # reorder explicitly afterwards
Defensive patterns

Strategy: validation

Validate before calling

for ax,(i,o) in enumerate(zip(in_spec, out_spec)):
    it = i if isinstance(i, tuple) else ((i,) if i else ())
    ot = o if isinstance(o, tuple) else ((o,) if o else ())
    assert o is None or it[:len(ot)] == ot, f'dim {ax}: out_spec not a prefix of in_spec'

Prevention

When it happens

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

Common situations: Trying to reorder gathered sub-axes or drop a non-leading sub-axis of a tuple-sharded dimension.

Related errors

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