jax-ml/jax · error · ValueError

vmap spmd_axis_name cannot appear in shard_map out_specs

Error message

vmap spmd_axis_name cannot appear in shard_map out_specs

What it means

Raised by the vmap batching rule for shard_map when the vmap spmd_axis_name also appears in shard_map's out_specs. The output batch dim is materialized via those SPMD axes, so naming them in out_specs double-specifies how to unshard outputs and is rejected.

Source

Thrown at jax/_src/shard_map.py:1606

  new_params = dict(mesh=mesh, in_specs=new_in_specs, check_vma=check_vma,
                    newly_manual_axes=newly_manual_axes, debug_info=debug_info)
  # TODO(yashkatariya): Remove remove_explicit_mesh_axis_names when vmap
  # mesh ctx is correctly set.
  with (core.set_current_trace(trace.parent_trace),
        core.remove_explicit_mesh_axis_names(trace.axis_data.explicit_mesh_axis)):
    out_vals = prim.bind(*in_vals, subfuns=(fun_batched,), **new_params)
  make_tracer = partial(batching.BatchTracer, trace,
                        source_info=source_info_util.current())
  out_vals, out_dims = out_vals.unpack_aux()
  return out_vals.map2(out_dims, make_tracer)
batching.BatchTrace.process_shard_map = _shard_map_batch

def _batch_out_specs(spmd_name, explicit_mesh_axis, dims, out_specs):
  if spmd_name is not None:
    used = {n for spec in out_specs for n in used_axis_names(spec)}
    if not config.disable_vmap_shmap_error.value and set(spmd_name) & used:
      raise ValueError("vmap spmd_axis_name cannot appear in shard_map out_specs")
    return [sp if d is None else pxla.batch_spec(sp, d, spmd_name)
            for sp, d in zip(out_specs, dims)]
  elif explicit_mesh_axis is not None:
    used = {n for spec in out_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 out_specs")
    return [sp if d is None else pxla.batch_spec(sp, d, None)
            for sp, d in zip(out_specs, dims)]
  else:
    return [sp if d is None else pxla.batch_spec(sp, d, None)
            for sp, d in zip(out_specs, dims)]


# Autodiff

def _shard_map_jvp(trace, shard_map_p, f, tracers, mesh, in_specs,
                   check_vma, newly_manual_axes, debug_info):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the spmd_axis_name axes from out_specs PartitionSpecs
2. Set jax.config.update('disable_vmap_shmap_error', True) if you knowingly want to bypass
3. Use distinct axis names for the vmap batch axis and the output sharding axes

Example fix

# before
shard_map(f, mesh=m, out_specs=PartitionSpec('i'))  # under vmap(spmd_axis_name='i')

# after
shard_map(f, mesh=m, out_specs=PartitionSpec(None))

spmd = {'i'}
used_out = {n for sp in out_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}
assert not (spmd & used_out), 'spmd_axis_name collides with shard_map out_specs'

def out_specs_disjoint_from_spmd(spmd_axis_name, out_specs) -> bool:
    spmd = {spmd_axis_name} if isinstance(spmd_axis_name, str) else set(spmd_axis_name)
    used = {n for sp in out_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 out_specs' in str(e):
        ...

When it happens

Trigger: jax.vmap(f, spmd_axis_name='i') wrapping a shard_map whose out_specs is PartitionSpec('i') (or otherwise mentions an axis in spmd_axis_name).

Common situations: Same class as the in_specs variant: SPMD-to-vmap migration or axis-name collisions after mesh refactors.

Related errors

• vmap spmd_axis_name cannot appear in shard_map in_specs
• vmapped away explicit mesh axis cannot appear in shard_map o
• 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

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