{"record":{"id":"7a448b23a24206fd","repo":"jax-ml/jax","slug":"shard-map-prefix-specs-argument-must-be-a-pytree","errorCode":null,"errorMessage":"shard_map {prefix}_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but:\n\n{msgs}\n\nCheck the {prefix}_specs values passed to shard_map.","messagePattern":"shard_map (.+?)_specs argument must be a pytree of `jax\\.sharding\\.PartitionSpec` instances, but:\n\n(.+?)\n\nCheck the (.+?)_specs values passed to shard_map\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":511,"sourceCode":"\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,\n    dyn_argnums: Sequence[int], in_specs_flat: Sequence[P],\n    xs: Sequence) -> None:\n  in_avals = map(core.shaped_abstractify, xs)\n  fail = [a if isinstance(p, P) and len(p) > a.ndim else no_fail\n          for p, a in zip(in_specs_flat, in_avals)]","sourceCodeStart":493,"sourceCodeEnd":529,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L493-L529","documentation":"Each leaf of in_specs/out_specs must be a jax.sharding.PartitionSpec instance. This TypeError lists every offending pytree path when a leaf is some other type (string, tuple, None, custom object).","triggerScenarios":"Passing raw tuples like in_specs=(('data','model'),) instead of PartitionSpec objects, or mixing strings/other objects into the spec pytree.","commonSituations":"Porting pmap-era code that used plain tuples; loading specs from JSON/YAML config and forgetting to convert; some tree leaves defaulting to 0/None.","solutions":["Wrap every spec leaf in PartitionSpec (P), e.g. P('data','model') instead of ('data','model')","Sanitize with jax.tree.map(lambda s: s if isinstance(s, P) else P(*s), specs)","Validate spec trees before calling shard_map"],"exampleFix":"// before\nshard_map(f, mesh, xs, in_specs=(('data',), None))\n// after\nfrom jax.sharding import PartitionSpec as P\nshard_map(f, mesh, xs, in_specs=(P('data'), P()))","handlingStrategy":"type-guard","validationCode":"from jax.sharding import PartitionSpec as P\nspecs = jax.tree.map(lambda s: s if isinstance(s, P) else P(*s) if isinstance(s, (tuple, list)) else s, specs)\nassert all(isinstance(s, P) for s in jax.tree.leaves(specs))","typeGuard":"def is_spec_pytree(specs) -> bool:\n    return all(isinstance(s, PartitionSpec) for s in jax.tree.leaves(specs))","tryCatchPattern":"try: shard_map(...) except TypeError as e: if 'pytree of' in str(e): coerce leaves to PartitionSpec and retry; else: raise","preventionTips":["Convert JSON/config specs to PartitionSpec at load time","Lint for bare tuples in spec trees"],"tags":["shard-map","partition-spec","typeerror","jax"],"backgroundTag":"invalid-sharding-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}