{"record":{"id":"e047fb92f25ec2a3","repo":"jax-ml/jax","slug":"vmap-spmd-axis-name-cannot-appear-in-shard-map-in","errorCode":null,"errorMessage":"vmap spmd_axis_name cannot appear in shard_map in_specs","messagePattern":"vmap spmd_axis_name cannot appear in shard_map in_specs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1559,"sourceCode":"eager_rules[core.empty_ref_p] = _ref_raise_valueerror\n\n# Batching\n\ndef used_axis_names(spec):\n  return _spec_to_mat(spec).vur\n\ndef _shard_map_batch(\n    trace: batching.BatchTrace, prim: core.Primitive, fun: Callable,\n    in_tracers: Sequence[batching.BatchTracer], mesh: Mesh,\n    in_specs, check_vma: bool, newly_manual_axes: frozenset,\n    debug_info) -> Sequence[batching.BatchTracer]:\n  in_vals, in_dims = unzip2(map(trace.to_batch_info, in_tracers))\n  spmd_axis_name = trace.axis_data.spmd_name\n  explicit_mesh_axis = trace.axis_data.explicit_mesh_axis\n  if spmd_axis_name is not None:\n    used = {n for spec in in_specs for n in used_axis_names(spec)}\n    if not config.disable_vmap_shmap_error.value and set(spmd_axis_name) & used:\n      raise ValueError(\"vmap spmd_axis_name cannot appear in shard_map in_specs\")\n    new_in_specs = [\n        sp if d is None else pxla.batch_spec(sp, d, spmd_axis_name)\n        for sp, d in zip(in_specs, in_dims)]\n    new_size = trace.axis_data.size // prod(mesh.shape[n] for n in spmd_axis_name)\n    new_axis_data = batching.AxisData(\n        trace.axis_data.name, new_size, trace.axis_data.spmd_name,\n        trace.axis_data.explicit_mesh_axis)\n  elif explicit_mesh_axis is not None:\n    used = {n for spec in in_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 in_specs\")\n    new_in_specs = [\n        sp if d is None else pxla.batch_spec(sp, d, None)\n        for sp, d in zip(in_specs, in_dims)]\n    new_axis_data = trace.axis_data\n  else:\n    new_in_specs = [sp if d is None else pxla.batch_spec(sp, d, None)","sourceCodeStart":1541,"sourceCodeEnd":1577,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1541-L1577","documentation":"Raised by JAX's vmap batching rule for shard_map when the axis name(s) passed as vmap's spmd_axis_name also appear in the shard_map in_specs PartitionSpecs. Because the batched dimension is meant to be mapped onto those SPMD mesh axes, referencing them again in in_specs is contradictory, so JAX rejects it (unless the escape hatch config.disable_vmap_shmap_error is set).","triggerScenarios":"Calling jax.vmap(f, spmd_axis_name='i') around a function that uses shard_map(..., in_specs=(PartitionSpec('i'),), mesh=...) where 'i' is in the spmd_axis_name set.","commonSituations":"Migrating manual pmap/pjit spmd code to vmap+shard_map; renaming mesh axes so the vmap spmd_axis_name accidentally collides with a spec axis name; using vmap over a model already wrapped in sharded inference code.","solutions":["Remove the spmd_axis_name mesh axis names from the shard_map in_specs (they are already consumed by vmap)","If the behavior is intended, set jax.config.update('disable_vmap_shmap_error', True) to downgrade this to a warning-free path","Double-check mesh axis naming: use distinct names for vmap-mapped axes vs axes you want sharded in specs","Restructure so the batched dimension is not simultaneously an SPMD axis (e.g. add a dedicated mesh axis for batching)"],"exampleFix":"# before\njax.vmap(f, spmd_axis_name='i')(x)  # f uses shard_map(..., in_specs=PartitionSpec('i',))\n\n# after\njax.vmap(f, spmd_axis_name='batch')(x)  # shard_map in_specs keeps or drops 'i' without collision","handlingStrategy":"validation","validationCode":"spmd = 'i'  # vmap spmd_axis_name\nused_in = {n for spec in in_specs_tree for n in (spec if isinstance(spec, tuple) else (spec,)) if n}\nassert not (set(spmd if isinstance(spmd, tuple) else (spmd,)) & used_in, 'spmd_axis_name collides with shard_map in_specs'","typeGuard":"def specs_disjoint_from_spmd(spmd_axis_name, in_specs) -> bool:\n    spmd = {spmd_axis_name} if isinstance(spmd_axis_name, str) else set(spmd_axis_name)\n    used = {n for sp in in_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 in_specs' in str(e):\n        # strip 'i' from in_specs and retry\n        ...","preventionTips":["Keep a strict naming convention: batch/spmd axes vs mesh sharding axes","Add a unit test asserting vmap-wrapped sharded functions run on tiny meshes","Lint PartitionSpecs for axis names not present in the mesh or colliding with spmd_axis_name"],"tags":["jax","vmap","shard-map","spmd","partition-spec"],"backgroundTag":"jax-sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}