jax-ml/jax · error · ValueError

vmapped away explicit mesh axis cannot appear in shard_map i

Error message

vmapped away explicit mesh axis cannot appear in shard_map in_specs

What it means

Raised when vmap removes (maps away) an explicit mesh axis that is still referenced in shard_map's in_specs. With explicit mesh axes, vmap consumes the axis, so its name can no longer be used to shard inputs inside shard_map.

Source

Thrown at jax/_src/shard_map.py:1570

    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)
                    for sp, d in zip(in_specs, in_dims)]
    new_axis_data = trace.axis_data

  def fun_batched(*args):
    ans_aux, out_dims = batching.batch_subtrace_2(
        fun, trace.tag, new_axis_data, tuple(in_dims), args)
    ans, out_specs = ans_aux.unpack_aux()
    new_out_specs = _batch_out_specs(spmd_axis_name, explicit_mesh_axis,
                                     out_dims, out_specs)
    return ans.with_aux(out_dims).with_aux(tuple(new_out_specs))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the vmapped-away explicit axis name from the in_specs PartitionSpecs
2. Move the batch dimension to a non-explicit (positional) axis or a dedicated mesh axis not consumed by vmap
3. Re-check which axes the vmap trace actually consumed by printing trace axis data / mesh

Example fix

# before
jax.vmap(f, axis_name='i')(x)  # 'i' is an explicit mesh axis, shard_map in_specs=PartitionSpec('i',)

# after
shard_map(..., in_specs=PartitionSpec(None,))  # let vmap own that dim

explicit_axes = {n for n, t in zip(mesh.axis_names, mesh.axis_types) if t is None}
used = {n for sp in in_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}
assert not (explicit_axes & used), 'vmapped explicit axis used in in_specs'

try:
    jax.vmap(f)(x)
except ValueError as e:
    if 'explicit mesh axis cannot appear in shard_map in_specs' in str(e):
        ...

When it happens

Trigger: Using an explicit abstract Mesh (jax.make_mesh(..., axis_types=(None,))) and vmap'ing over an axis that appears in the shard_map in_specs PartitionSpec.

Common situations: Adopting the new explicit-mesh API and mixing vmap over explicit axes with sharding specs that still name those axes.

Related errors

• vmapped away explicit mesh axis cannot appear in shard_map o
• vmap spmd_axis_name cannot appear in shard_map in_specs
• vmap spmd_axis_name cannot appear in shard_map out_specs
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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