jax-ml/jax · error · ValueError

vmapped away explicit mesh axis cannot appear in shard_map o

Error message

vmapped away explicit mesh axis cannot appear in shard_map out_specs

What it means

Raised when vmap maps away an explicit mesh axis that is still named in shard_map's out_specs. Since vmap consumes the axis to produce the batch dimension, it cannot simultaneously define how shard_map outputs are unsharded.

Source

Thrown at jax/_src/shard_map.py:1612

        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):
  debug_info = debug_info.with_unknown_names()
  primals, tangents = unzip2(map(trace.to_primal_tangent_pair, tracers))
  which_nz = [     type(t) is not ad.Zero           for t in tangents]
  tangents = [t if type(t) is not ad.Zero else None for t in tangents]
  args, in_zeros_tree = tree_flatten((primals, tangents))
  tangent_in_specs = [sp.to_tangent_spec() for sp, nz in zip(in_specs, which_nz) if nz]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop the vmapped-away explicit axis from out_specs
Restructure so vmap's batch dim is not an explicit mesh axis used in output specs

Example fix

# before
shard_map(f, mesh=m, out_specs=PartitionSpec('i'))  # 'i' explicit and vmapped

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

Defensive patterns

Strategy: validation

Validate before calling

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

Try / catch

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

Prevention

Track which explicit axes are already consumed by vmap at each nesting level

When it happens

Trigger: vmap over an explicit mesh axis combined with shard_map out_specs that reference that axis name.

Common situations: Explicit-mesh (axis_types) codebases adding vmap on top of existing sharded functions.

Related errors

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