{"record":{"id":"fddbdc52c2d0b78b","repo":"jax-ml/jax","slug":"vmapped-away-explicit-mesh-axis-cannot-appear-in-s-fddbdc","errorCode":null,"errorMessage":"vmapped away explicit mesh axis cannot appear in shard_map out_specs","messagePattern":"vmapped away explicit mesh axis cannot appear in shard_map out_specs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1612,"sourceCode":"        core.remove_explicit_mesh_axis_names(trace.axis_data.explicit_mesh_axis)):\n    out_vals = prim.bind(*in_vals, subfuns=(fun_batched,), **new_params)\n  make_tracer = partial(batching.BatchTracer, trace,\n                        source_info=source_info_util.current())\n  out_vals, out_dims = out_vals.unpack_aux()\n  return out_vals.map2(out_dims, make_tracer)\nbatching.BatchTrace.process_shard_map = _shard_map_batch\n\ndef _batch_out_specs(spmd_name, explicit_mesh_axis, dims, out_specs):\n  if spmd_name is not None:\n    used = {n for spec in out_specs for n in used_axis_names(spec)}\n    if not config.disable_vmap_shmap_error.value and set(spmd_name) & used:\n      raise ValueError(\"vmap spmd_axis_name cannot appear in shard_map out_specs\")\n    return [sp if d is None else pxla.batch_spec(sp, d, spmd_name)\n            for sp, d in zip(out_specs, dims)]\n  elif explicit_mesh_axis is not None:\n    used = {n for spec in out_specs for n in used_axis_names(spec)}\n    if set(explicit_mesh_axis) & used:\n      raise ValueError(\"vmapped away explicit mesh axis cannot appear in \"\n                       \"shard_map out_specs\")\n    return [sp if d is None else pxla.batch_spec(sp, d, None)\n            for sp, d in zip(out_specs, dims)]\n  else:\n    return [sp if d is None else pxla.batch_spec(sp, d, None)\n            for sp, d in zip(out_specs, dims)]\n\n\n# Autodiff\n\ndef _shard_map_jvp(trace, shard_map_p, f, tracers, mesh, in_specs,\n                   check_vma, newly_manual_axes, debug_info):\n  debug_info = debug_info.with_unknown_names()\n  primals, tangents = unzip2(map(trace.to_primal_tangent_pair, tracers))\n  which_nz = [     type(t) is not ad.Zero           for t in tangents]\n  tangents = [t if type(t) is not ad.Zero else None for t in tangents]\n  args, in_zeros_tree = tree_flatten((primals, tangents))\n  tangent_in_specs = [sp.to_tangent_spec() for sp, nz in zip(in_specs, which_nz) if nz]","sourceCodeStart":1594,"sourceCodeEnd":1630,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1594-L1630","documentation":"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.","triggerScenarios":"vmap over an explicit mesh axis combined with shard_map out_specs that reference that axis name.","commonSituations":"Explicit-mesh (axis_types) codebases adding vmap on top of existing sharded functions.","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"],"exampleFix":"# before\nshard_map(f, mesh=m, out_specs=PartitionSpec('i'))  # 'i' explicit and vmapped\n\n# after\nshard_map(f, mesh=m, out_specs=PartitionSpec(None))","handlingStrategy":"validation","validationCode":"explicit_axes = {n for n, t in zip(mesh.axis_names, mesh.axis_types) if t is None}\nused = {n for sp in out_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}\nassert not (explicit_axes & used), 'vmapped explicit axis used in out_specs'","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(f)(x)\nexcept ValueError as e:\n    if 'explicit mesh axis cannot appear in shard_map out_specs' in str(e):\n        ...","preventionTips":["Track which explicit axes are already consumed by vmap at each nesting level"],"tags":["jax","vmap","shard-map","explicit-mesh","out-specs"],"backgroundTag":"jax-sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}