jax-ml/jax · error · ValueError

jax.shard_map requires axis_names={axis_names} to be a subse

Error message

jax.shard_map requires axis_names={axis_names} to be a subset of mesh.axis_names={mesh_axis_names_wo_vmap}

What it means

shard_map requires every name in axis_names to exist on the mesh (specifically the mesh's non-vmap axis names). If a name is not a mesh axis — a typo, or an axis from a different mesh — this ValueError lists both sets so you can see the discrepancy.

Source

Thrown at jax/_src/shard_map.py:398

                    "`jax.sharding.AbstractMesh` instance for its "
                    f"second argument, but got {mesh} of type {type(mesh)}.")
  if mesh.empty:
    raise ValueError(f"shard_map requires a non-empty mesh. Got {mesh}")

  mesh_axis_names_wo_vmap = (
      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names
  )

  if not isinstance(axis_names, (frozenset, set)):
    raise TypeError(
        "`axis_names` argument of shard_map should be of type `frozenset` or"
        f" `set`. Got type: {type(axis_names)}")
  if isinstance(axis_names, set):
    axis_names = frozenset(axis_names)
  if not axis_names:
    axis_names = mesh_axis_names_wo_vmap
  if not axis_names.issubset(mesh_axis_names_wo_vmap):
    raise ValueError(
        f"jax.shard_map requires axis_names={axis_names} to be a subset of "
        f"mesh.axis_names={mesh_axis_names_wo_vmap}")

  if (in_specs is Infer and
      not all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):
    axis_types = ', '.join(str(mesh._name_to_type[a]) for a in axis_names)
    if _smap:
      msg = (f"in_axes was not specified when axis_name={axis_names} was of"
             f" type {axis_types}")
    else:
      msg = ("shard_map in_specs argument must be a pytree of"
             " `jax.sharding.PartitionSpec` instances, but it was `None` when"
             f" {axis_names=} are of type {axis_types}")
    raise TypeError(msg)

  if in_specs is not Infer and in_specs is not None:
    _check_specs(SpecErrorType.input, in_specs, axis_names)
    _check_unreduced(SpecErrorType.input, mesh, axis_names, in_specs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align names: use only axis names present in mesh.axis_names, e.g. axis_names=frozenset(mesh.axis_names)
2. Fix typos or update the Mesh definition so its axis_names include the names you reference

Example fix

// before
mesh = jax.sharding.Mesh(devices, axis_names=('i',))
jax.shard_map(f, mesh=mesh, axis_names=frozenset({'data'}), ...)

// after
jax.shard_map(f, mesh=mesh, axis_names=frozenset({'i'}), ...)

mesh_axis_names = set(mesh.axis_names)
assert axis_names <= mesh_axis_names, f'{axis_names - mesh_axis_names} not in mesh'

def names_in_mesh(ns, mesh) -> bool:
    return set(ns) <= set(mesh.axis_names)

When it happens

Trigger: Calling shard_map with axis_names={'data'} when the mesh was created with axis_names=('i',) — or reusing specs/names written for one mesh layout against a differently-named mesh.

Common situations: Renaming mesh axes during refactors without updating shard_map call sites; copy-pasted model code assuming a standard ('data','model') mesh while the local mesh uses different names; subsetting names after mesh reconfiguration.

Related errors

• The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
• The context mesh {ctx_mesh} should match the mesh passed to
• shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.
• shard_map requires a non-empty mesh. Got {mesh}
• Mapped away dimension of inputs passed to vmap should be sha

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