jax-ml/jax · error · ValueError

top_level_all_gather works when all mesh axes of context mes

Error message

top_level_all_gather works when all mesh axes of context mesh are explicit. Got {get_abstract_mesh()}

What it means

top_level_all_gather only supports context meshes whose axes are all explicit (created with axis_types where every axis is named/explicit). If the ambient abstract mesh has positional/implicit axes, the API refuses to run.

Source

Thrown at jax/_src/shard_map.py:2248

            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,
      tupled_args=True)
  x_avals_flat = [core.typeof(x) for x in x_flat]
  out_flat = [_top_level_ag(x, aval, sh, multi_dim)
              for x, aval, sh in zip(x_flat, x_avals_flat, out_sharding_flat)]
  return tree_unflatten(treedef, out_flat)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Enter an explicit mesh context before calling (e.g. `with mesh: top_level_all_gather(...)`), where mesh was built with all-explicit axes
2. Rebuild the mesh with jax.make_mesh(shape, names, axis_types=(...)) with no None entries
3. Or use shard_map/pjit-based all_gather instead of top_level_all_gather

Example fix

# before
top_level_all_gather(x, out_sh)  # context mesh has implicit axes

# after
with explicit_mesh:
  top_level_all_gather(x, out_sh)

from jax.experimental.mesh_utils import get_abstract_mesh  # or jax._src.mesh
m = get_abstract_mesh()
assert m.are_all_axes_explicit, f'need all-explicit mesh, got {m}'

When it happens

Trigger: Calling top_level_all_gather while the context mesh (set via mesh context manager or with_mesh) contains non-explicit axes (axis_types=None entries).

Common situations: Using new-style jax.make_mesh(..., axis_types=...) partially explicit, or forgetting to enter an explicit mesh context before the call.

Related errors

• all_gather_reduced only accepts inputs that are varying. Got
• all_gather_reduced is a Varying -> Reduced collective. This
• all_gather_reduced's input cannot be reduced across the axis
• Please file an issue at https://github.com/jax-ml/jax/issues
• vmapped away explicit mesh axis cannot appear in shard_map i

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