{"record":{"id":"1517937b3e01f282","repo":"jax-ml/jax","slug":"jax-shard-map-requires-axis-names-axis-names-to","errorCode":null,"errorMessage":"jax.shard_map requires axis_names={axis_names} to be a subset of mesh.axis_names={mesh_axis_names_wo_vmap}","messagePattern":"jax\\.shard_map requires axis_names=(.+?) to be a subset of mesh\\.axis_names=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":398,"sourceCode":"                    \"`jax.sharding.AbstractMesh` instance for its \"\n                    f\"second argument, but got {mesh} of type {type(mesh)}.\")\n  if mesh.empty:\n    raise ValueError(f\"shard_map requires a non-empty mesh. Got {mesh}\")\n\n  mesh_axis_names_wo_vmap = (\n      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names\n  )\n\n  if not isinstance(axis_names, (frozenset, set)):\n    raise TypeError(\n        \"`axis_names` argument of shard_map should be of type `frozenset` or\"\n        f\" `set`. Got type: {type(axis_names)}\")\n  if isinstance(axis_names, set):\n    axis_names = frozenset(axis_names)\n  if not axis_names:\n    axis_names = mesh_axis_names_wo_vmap\n  if not axis_names.issubset(mesh_axis_names_wo_vmap):\n    raise ValueError(\n        f\"jax.shard_map requires axis_names={axis_names} to be a subset of \"\n        f\"mesh.axis_names={mesh_axis_names_wo_vmap}\")\n\n  if (in_specs is Infer and\n      not all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):\n    axis_types = ', '.join(str(mesh._name_to_type[a]) for a in axis_names)\n    if _smap:\n      msg = (f\"in_axes was not specified when axis_name={axis_names} was of\"\n             f\" type {axis_types}\")\n    else:\n      msg = (\"shard_map in_specs argument must be a pytree of\"\n             \" `jax.sharding.PartitionSpec` instances, but it was `None` when\"\n             f\" {axis_names=} are of type {axis_types}\")\n    raise TypeError(msg)\n\n  if in_specs is not Infer and in_specs is not None:\n    _check_specs(SpecErrorType.input, in_specs, axis_names)\n    _check_unreduced(SpecErrorType.input, mesh, axis_names, in_specs)","sourceCodeStart":380,"sourceCodeEnd":416,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L380-L416","documentation":"shard_map requires every name in axis_names to exist on the mesh (specifically the mesh's non-vmap axis names). If a name is not a mesh axis — a typo, or an axis from a different mesh — this ValueError lists both sets so you can see the discrepancy.","triggerScenarios":"Calling shard_map with axis_names={'data'} when the mesh was created with axis_names=('i',) — or reusing specs/names written for one mesh layout against a differently-named mesh.","commonSituations":"Renaming mesh axes during refactors without updating shard_map call sites; copy-pasted model code assuming a standard ('data','model') mesh while the local mesh uses different names; subsetting names after mesh reconfiguration.","solutions":["Align names: use only axis names present in mesh.axis_names, e.g. axis_names=frozenset(mesh.axis_names)","Fix typos or update the Mesh definition so its axis_names include the names you reference"],"exampleFix":"// before\nmesh = jax.sharding.Mesh(devices, axis_names=('i',))\njax.shard_map(f, mesh=mesh, axis_names=frozenset({'data'}), ...)\n\n// after\njax.shard_map(f, mesh=mesh, axis_names=frozenset({'i'}), ...)","handlingStrategy":"validation","validationCode":"mesh_axis_names = set(mesh.axis_names)\nassert axis_names <= mesh_axis_names, f'{axis_names - mesh_axis_names} not in mesh'","typeGuard":"def names_in_mesh(ns, mesh) -> bool:\n    return set(ns) <= set(mesh.axis_names)","tryCatchPattern":null,"preventionTips":["Derive axis_names from mesh.axis_names instead of hardcoding","Centralize mesh construction so names stay consistent","Add a smoke test asserting spec names ⊆ mesh names"],"tags":["jax","shard-map","axis-names","mesh","distributed"],"backgroundTag":"axis-name-not-in-mesh","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}