jax-ml/jax · error · ValueError

out_sharding passed to top_level_all_gather cannot be {out_s

Error message

out_sharding passed to top_level_all_gather cannot be {out_sh_}. It should be a PartitionSpec or NamedSharding.

What it means

top_level_all_gather canonicalizes its out_sharding argument; if canonicalization returns None the value is not a PartitionSpec or NamedSharding (e.g. a raw string, GSPMDSharding, or wrong object), and JAX raises this error.

Source

Thrown at jax/_src/shard_map.py:2203

      newly_manual_axes=newly_manual_axes, debug_info=debug_info,
      check_vma=check_vma)
  out_vals, ref_vals = split_list(out_and_ref_vals, [len(jaxpr.outvars)])
  ref_vals_ = iter(ref_vals)
  new_invals = [next(ref_vals_) if isinstance(a, AbstractRef) else None
                for a in ctx.in_avals]
  assert next(ref_vals_, None) is None
  return new_invals, out_vals

def _repspec(aval):
  return aval.nospec(empty_abstract_mesh, False, ())

# ----------------------- top level collectives --------------------------------

def _top_level_ag(x, aval, out_sh_, multi_dim):
  assert aval.sharding.mesh.are_all_axes_explicit, aval.sharding.mesh
  out_sh = canonicalize_sharding(out_sh_, "top_level_all_gather")
  if out_sh is None:
    raise ValueError(
        f'out_sharding passed to top_level_all_gather cannot be {out_sh_}. It'
        ' should be a PartitionSpec or NamedSharding.')
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a NamedSharding(mesh, PartitionSpec(...)) or a plain PartitionSpec as out_sharding
2. If you have an HloSharding, convert it to a NamedSharding first
3. Check for None/typo'd variables being passed through

Example fix

# before
top_level_all_gather(x, out_sharding='data')

# after
from jax.sharding import NamedSharding, PartitionSpec as P
top_level_all_gather(x, out_sharding=NamedSharding(mesh, P('data')))

from jax.sharding import PartitionSpec, NamedSharding
assert isinstance(out_sharding, (PartitionSpec, NamedSharding)), f'bad out_sharding type: {type(out_sharding)}'

def is_valid_out_sharding(s) -> bool:
    from jax.sharding import PartitionSpec, NamedSharding
    return isinstance(s, (PartitionSpec, NamedSharding))

When it happens

Trigger: Calling jax.experimental.top_level_all_gather (shard_map module) with out_sharding that is not a PartitionSpec/NamedSharding, such as a string spec or an OpSharding proto.

Common situations: Passing a PjitSharding/GSPMDSharding or hand-built object where a user-facing PartitionSpec is expected; API confusion with pjit's in_shardings which accept more types.

Related errors

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