jax-ml/jax · error · ValueError

Input sharding mesh {aval.sharding.mesh} should be equal to

Error message

Input sharding mesh {aval.sharding.mesh} should be equal to out_sharding mesh {out_sh.mesh}

What it means

top_level_all_gather requires the mesh implied by the input array's sharding to equal the mesh of the out_sharding. Collectives are emitted on the input's mesh, so a mismatched out_sharding mesh (different devices or axis names) is rejected.

Source

Thrown at jax/_src/shard_map.py:2207

  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:
        continue
      if not multi_dim and count > 0:
        raise ValueError(
            "multiple dimensions cannot be all_gathered since multi_dim=False"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build out_sharding from the same Mesh object used for the input's sharding (or the context mesh)
2. Ensure mesh axis names match exactly — meshes differing only in names still compare unequal
3. Shard the input on the same mesh you intend to gather on before calling

Example fix

# before
out = top_level_all_gather(x, NamedSharding(other_mesh, P('data')))

# after
out = top_level_all_gather(x, NamedSharding(x_mesh, P('data')))  # x_mesh = mesh x is sharded on

in_mesh = x.sharding.mesh if hasattr(x, 'sharding') else get_abstract_mesh()
assert in_mesh == out_sharding.mesh if isinstance(out_sharding, NamedSharding) else True, 'mesh mismatch'

def meshes_match(x, out_sharding) -> bool:
    return isinstance(out_sharding, NamedSharding) and x.sharding.mesh == out_sharding.mesh

When it happens

Trigger: Input array sharded on mesh A (e.g. via jit in_shardings) but out_sharding built from a different mesh B, or meshes with same shape but different axis names.

Common situations: Creating multiple Mesh objects in one process; refactoring context meshes so names differ; mixing jax.make_mesh results across modules.

Related errors

• Mesh must be provided for shard_map with checkify.
• Found an unbound axis name: {name}. To fix this, please call
• 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

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