{"record":{"id":"a79b868c93050326","repo":"jax-ml/jax","slug":"vmap-spmd-axis-name-cannot-appear-in-shard-map-out","errorCode":null,"errorMessage":"vmap spmd_axis_name cannot appear in shard_map out_specs","messagePattern":"vmap spmd_axis_name cannot appear in shard_map out_specs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1606,"sourceCode":"\n  new_params = dict(mesh=mesh, in_specs=new_in_specs, check_vma=check_vma,\n                    newly_manual_axes=newly_manual_axes, debug_info=debug_info)\n  # TODO(yashkatariya): Remove remove_explicit_mesh_axis_names when vmap\n  # mesh ctx is correctly set.\n  with (core.set_current_trace(trace.parent_trace),\n        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):","sourceCodeStart":1588,"sourceCodeEnd":1624,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1588-L1624","documentation":"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.","triggerScenarios":"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).","commonSituations":"Same class as the in_specs variant: SPMD-to-vmap migration or axis-name collisions after mesh refactors.","solutions":["Remove the spmd_axis_name axes from out_specs PartitionSpecs","Set jax.config.update('disable_vmap_shmap_error', True) if you knowingly want to bypass","Use distinct axis names for the vmap batch axis and the output sharding axes"],"exampleFix":"# before\nshard_map(f, mesh=m, out_specs=PartitionSpec('i'))  # under vmap(spmd_axis_name='i')\n\n# after\nshard_map(f, mesh=m, out_specs=PartitionSpec(None))","handlingStrategy":"validation","validationCode":"spmd = {'i'}\nused_out = {n for sp in out_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}\nassert not (spmd & used_out), 'spmd_axis_name collides with shard_map out_specs'","typeGuard":"def out_specs_disjoint_from_spmd(spmd_axis_name, out_specs) -> bool:\n    spmd = {spmd_axis_name} if isinstance(spmd_axis_name, str) else set(spmd_axis_name)\n    used = {n for sp in out_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}\n    return not (spmd & used)","tryCatchPattern":"try:\n    jax.vmap(f, spmd_axis_name='i')(x)\nexcept ValueError as e:\n    if 'spmd_axis_name cannot appear in shard_map out_specs' in str(e):\n        ...","preventionTips":["Centralize spec construction so in/out specs derive from one axis-naming helper","Avoid reusing mesh axis names as vmap axis_name/spmd_axis_name"],"tags":["jax","vmap","shard-map","spmd","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"}