{"record":{"id":"98852e2ef0922f1a","repo":"jax-ml/jax","slug":"shard-map-prefix-specs-argument-must-refer-to-an","errorCode":null,"errorMessage":"shard_map {prefix}_specs argument must refer to an axis marked as manual ({manual_axes}), but:\n\n{msgs}\n\nCheck the {prefix}_specs values passed to shard_map.","messagePattern":"shard_map (.+?)_specs argument must refer to an axis marked as manual \\((.+?)\\), but:\n\n(.+?)\n\nCheck the (.+?)_specs values passed to shard_map\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":506,"sourceCode":"      names = (names,) if not isinstance(names, tuple) else names\n      for name in names:\n        if name is not None and name not in manual_axes:\n          return False\n    return True\n\n  if all(check_spec(p) for p in tree_leaves(specs)):\n    return\n  prefix = 'in' if error_type == SpecErrorType.input else 'out'\n  msgs = [f\"  {prefix}_specs{keystr(key)} is {x} of type {type(x).__name__}, \"\n          for key, x in generate_key_paths(specs) if not isinstance(x, P)]\n  if not msgs:\n    for key, p in generate_key_paths(specs):\n      for names in p:\n        names = (names,) if not isinstance(names, tuple) else names\n        for name in names:\n          if name is not None and name not in manual_axes:\n            msgs.append(f\"  {prefix}_specs{keystr(key)} refers to {repr(name)}\")\n    raise ValueError(\n        f\"shard_map {prefix}_specs argument must refer to an axis \"\n        f\"marked as manual ({manual_axes}), but:\\n\\n\"\n        + '\\n\\n'.join(msgs) + '\\n\\n'\n        f\"Check the {prefix}_specs values passed to shard_map.\")\n  raise TypeError(\n      f\"shard_map {prefix}_specs argument must be a pytree of \"\n      f\"`jax.sharding.PartitionSpec` instances, but:\\n\\n\"\n      + '\\n\\n'.join(msgs) + '\\n\\n'\n      f\"Check the {prefix}_specs values passed to shard_map.\")\n\nclass NoFail:\n  def __repr__(self):\n    return \"NoFail()\"\n\nno_fail = NoFail()\n\ndef _check_specs_vs_args(\n    f: Callable, mesh: Mesh | AbstractMesh, in_tree: PyTreeDef, in_specs: Specs,","sourceCodeStart":488,"sourceCodeEnd":524,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L488-L524","documentation":"In manual/partial-manual shard_map mode, specs may only reference mesh axes that are marked manual (via manual_axes). This ValueError lists each spec path that references a non-manual axis name.","triggerScenarios":"Calling shard_map(..., manual_axes=('m',)) (or checking a spec in a partially-manual context) with in_specs/out_specs naming an axis not included in manual_axes, e.g. out_specs=P('data') when 'data' is not manual.","commonSituations":"Migrating code to manual sharding mode and forgetting to update specs; using axis names from the global mesh that were not declared as manual axes.","solutions":["Add the referenced axis name to manual_axes in the shard_map call","Change the spec to only use axis names already in manual_axes (or use None to replicate)","Re-check which axes the surrounding computation actually treats as manual"],"exampleFix":"// before\nshard_map(f, mesh, manual_axes=('m',), in_specs=P('m','x'))\n// after\nshard_map(f, mesh, manual_axes=('m','x'), in_specs=P('m','x'))","handlingStrategy":"validation","validationCode":"def specs_within_manual(specs, manual_axes):\n    bad = [n for p in jax.tree.leaves(specs) if isinstance(p, P)\n           for names in p for n in ((names,) if not isinstance(names, tuple) else names)\n           if n is not None and n not in manual_axes]\n    return not bad","typeGuard":"def manual_safe(specs, manual_axes) -> bool:\n    ok = lambda p: all(n is None or n in manual_axes for n in (x for part in p for x in (part if isinstance(part, tuple) else (part,))))\n    return all(ok(p) for p in jax.tree.leaves(specs) if isinstance(p, P))","tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'marked as manual' in str(e): parse listed names and add to manual_axes; else: raise","preventionTips":["Derive manual_axes from the axes actually named in specs","Keep spec axis-name sets in one constants module"],"tags":["shard-map","manual-axes","partition-spec","jax"],"backgroundTag":"axis-name-not-in-scope","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}