{"record":{"id":"6058a2491b93c16b","repo":"jax-ml/jax","slug":"vmapped-away-explicit-mesh-axis-cannot-appear-in-s","errorCode":null,"errorMessage":"vmapped away explicit mesh axis cannot appear in shard_map in_specs","messagePattern":"vmapped away explicit mesh axis cannot appear in shard_map in_specs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1570,"sourceCode":"    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)\n                    for sp, d in zip(in_specs, in_dims)]\n    new_axis_data = trace.axis_data\n\n  def fun_batched(*args):\n    ans_aux, out_dims = batching.batch_subtrace_2(\n        fun, trace.tag, new_axis_data, tuple(in_dims), args)\n    ans, out_specs = ans_aux.unpack_aux()\n    new_out_specs = _batch_out_specs(spmd_axis_name, explicit_mesh_axis,\n                                     out_dims, out_specs)\n    return ans.with_aux(out_dims).with_aux(tuple(new_out_specs))\n","sourceCodeStart":1552,"sourceCodeEnd":1588,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1552-L1588","documentation":"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.","triggerScenarios":"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.","commonSituations":"Adopting the new explicit-mesh API and mixing vmap over explicit axes with sharding specs that still name those axes.","solutions":["Remove the vmapped-away explicit axis name from the in_specs PartitionSpecs","Move the batch dimension to a non-explicit (positional) axis or a dedicated mesh axis not consumed by vmap","Re-check which axes the vmap trace actually consumed by printing trace axis data / mesh"],"exampleFix":"# before\njax.vmap(f, axis_name='i')(x)  # 'i' is an explicit mesh axis, shard_map in_specs=PartitionSpec('i',)\n\n# after\nshard_map(..., in_specs=PartitionSpec(None,))  # let vmap own that dim","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 in_specs for n in (sp if isinstance(sp, tuple) else (sp,)) if n}\nassert not (explicit_axes & used), 'vmapped explicit axis used in in_specs'","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(f)(x)\nexcept ValueError as e:\n    if 'explicit mesh axis cannot appear in shard_map in_specs' in str(e):\n        ...","preventionTips":["Document which explicit axes each vmap consumes","Prefer non-explicit axes for pure batch dims","Test vmap+shard_map combos in CI on a 1-process mesh"],"tags":["jax","vmap","shard-map","explicit-mesh"],"backgroundTag":"jax-sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}