jax-ml/jax · error · ValueError

vmap spmd_axis_name cannot appear in shard_map in_specs

Error message

vmap spmd_axis_name cannot appear in shard_map in_specs

What it means

Raised by JAX's vmap batching rule for shard_map when the axis name(s) passed as vmap's spmd_axis_name also appear in the shard_map in_specs PartitionSpecs. Because the batched dimension is meant to be mapped onto those SPMD mesh axes, referencing them again in in_specs is contradictory, so JAX rejects it (unless the escape hatch config.disable_vmap_shmap_error is set).

Source

Thrown at jax/_src/shard_map.py:1559

eager_rules[core.empty_ref_p] = _ref_raise_valueerror

# Batching

def used_axis_names(spec):
  return _spec_to_mat(spec).vur

def _shard_map_batch(
    trace: batching.BatchTrace, prim: core.Primitive, fun: Callable,
    in_tracers: Sequence[batching.BatchTracer], mesh: Mesh,
    in_specs, check_vma: bool, newly_manual_axes: frozenset,
    debug_info) -> Sequence[batching.BatchTracer]:
  in_vals, in_dims = unzip2(map(trace.to_batch_info, in_tracers))
  spmd_axis_name = trace.axis_data.spmd_name
  explicit_mesh_axis = trace.axis_data.explicit_mesh_axis
  if spmd_axis_name is not None:
    used = {n for spec in in_specs for n in used_axis_names(spec)}
    if not config.disable_vmap_shmap_error.value and set(spmd_axis_name) & used:
      raise ValueError("vmap spmd_axis_name cannot appear in shard_map in_specs")
    new_in_specs = [
        sp if d is None else pxla.batch_spec(sp, d, spmd_axis_name)
        for sp, d in zip(in_specs, in_dims)]
    new_size = trace.axis_data.size // prod(mesh.shape[n] for n in spmd_axis_name)
    new_axis_data = batching.AxisData(
        trace.axis_data.name, new_size, trace.axis_data.spmd_name,
        trace.axis_data.explicit_mesh_axis)
  elif explicit_mesh_axis is not None:
    used = {n for spec in in_specs for n in used_axis_names(spec)}
    if set(explicit_mesh_axis) & used:
      raise ValueError("vmapped away explicit mesh axis cannot appear in "
                       "shard_map in_specs")
    new_in_specs = [
        sp if d is None else pxla.batch_spec(sp, d, None)
        for sp, d in zip(in_specs, in_dims)]
    new_axis_data = trace.axis_data
  else:
    new_in_specs = [sp if d is None else pxla.batch_spec(sp, d, None)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the spmd_axis_name mesh axis names from the shard_map in_specs (they are already consumed by vmap)
2. If the behavior is intended, set jax.config.update('disable_vmap_shmap_error', True) to downgrade this to a warning-free path
3. Double-check mesh axis naming: use distinct names for vmap-mapped axes vs axes you want sharded in specs
4. Restructure so the batched dimension is not simultaneously an SPMD axis (e.g. add a dedicated mesh axis for batching)

Example fix

# before
jax.vmap(f, spmd_axis_name='i')(x)  # f uses shard_map(..., in_specs=PartitionSpec('i',))

# after
jax.vmap(f, spmd_axis_name='batch')(x)  # shard_map in_specs keeps or drops 'i' without collision

spmd = 'i'  # vmap spmd_axis_name
used_in = {n for spec in in_specs_tree for n in (spec if isinstance(spec, tuple) else (spec,)) if n}
assert not (set(spmd if isinstance(spmd, tuple) else (spmd,)) & used_in, 'spmd_axis_name collides with shard_map in_specs'

def specs_disjoint_from_spmd(spmd_axis_name, in_specs) -> bool:
    spmd = {spmd_axis_name} if isinstance(spmd_axis_name, str) else set(spmd_axis_name)
    used = {n for sp in in_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}
    return not (spmd & used)

try:
    jax.vmap(f, spmd_axis_name='i')(x)
except ValueError as e:
    if 'spmd_axis_name cannot appear in shard_map in_specs' in str(e):
        # strip 'i' from in_specs and retry
        ...

When it happens

Trigger: Calling jax.vmap(f, spmd_axis_name='i') around a function that uses shard_map(..., in_specs=(PartitionSpec('i'),), mesh=...) where 'i' is in the spmd_axis_name set.

Common situations: Migrating manual pmap/pjit spmd code to vmap+shard_map; renaming mesh axes so the vmap spmd_axis_name accidentally collides with a spec axis name; using vmap over a model already wrapped in sharded inference code.

Related errors

• vmap spmd_axis_name cannot appear in shard_map out_specs
• Primitive {prim_name} requires varying manual axes to match,
• Scan carry input and output got mismatched varying manual ax
• {name} cannot accept args with unreduced_kind={mat.unreduced
• in_specs passed to shard_map: {s} does not match the specs o

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