{"record":{"id":"f60a83fee1b9a9a5","repo":"jax-ml/jax","slug":"shard-map-in-specs-argument-must-be-a-pytree-of-j","errorCode":null,"errorMessage":"shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but it was None.\nInstead of `in_specs=None`, did you mean `in_specs=P()`, where `P = jax.sharding.PartitionSpec`?","messagePattern":"shard_map in_specs argument must be a pytree of `jax\\.sharding\\.PartitionSpec` instances, but it was None\\.\nInstead of `in_specs=None`, did you mean `in_specs=P\\(\\)`, where `P = jax\\.sharding\\.PartitionSpec`\\?","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":476,"sourceCode":"          \" shard_map is in full manual mode. Got mesh axis names\"\n          f\" {mesh.axis_names}, manual_axes: {manual_axes}, specs: {s}. Please\"\n          \" file a bug at https://github.com/jax-ml/jax/issues.\")\n    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced):\n      raise ValueError(\n          f\"unreduced in {prefix}_specs {s} can only be used when the mesh\"\n          \" passed to shard_map contains axis names all of type `Explicit`.\"\n          f\" Got mesh {mesh}\")\n    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.reduced):\n      raise ValueError(\n          f\"reduced in {prefix}_specs {s} can only be used when the mesh\"\n          \" passed to shard_map contains axis names all of type `Explicit`.\"\n          f\" Got mesh {mesh}\")\n\n\ndef _check_specs(error_type: SpecErrorType, specs: Any, manual_axes) -> None:\n  from jax._src.hijax import HiPspec\n  if error_type == SpecErrorType.input and specs is None:\n    raise TypeError(\n        \"shard_map in_specs argument must be a pytree of \"\n        \"`jax.sharding.PartitionSpec` instances, but it was None.\\n\"\n        \"Instead of `in_specs=None`, did you mean `in_specs=P()`, \"\n        \"where `P = jax.sharding.PartitionSpec`?\")\n\n  def check_spec(p):\n    if isinstance(p, HiPspec):\n      return True  # TODO(mattjj,yashkatariya): add user validation method\n    if not isinstance(p, PartitionSpec):\n      return False\n    for names in p.partitions:\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          return False\n    return True\n\n  if all(check_spec(p) for p in tree_leaves(specs)):","sourceCodeStart":458,"sourceCodeEnd":494,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L458-L494","documentation":"shard_map validates that in_specs is a pytree of PartitionSpec instances; passing the Python value None (e.g. as a placeholder or default) is a TypeError because None is not a valid spec. The empty spec P() (which replicates) is almost always what was meant.","triggerScenarios":"Calling shard_map(f, mesh, in_specs=None) or having a None leaf inside a pytree passed as in_specs, often from a defaulted function parameter or a config-driven spec tree.","commonSituations":"Functions with `in_specs=None` default arguments; dataclass/config fields that default to None and are forwarded into shard_map.","solutions":["Replace None with P() (from jax.sharding.PartitionSpec) to mean 'replicated'","Use tree_map(lambda x: P() if x is None else x, in_specs) to sanitize pytrees","Default function parameters to P() instead of None"],"exampleFix":"// before\nshard_map(f, mesh, xs, in_specs=None)\n// after\nfrom jax.sharding import PartitionSpec as P\nshard_map(f, mesh, xs, in_specs=P())","handlingStrategy":"type-guard","validationCode":"import jax\nfrom jax.sharding import PartitionSpec as P\nspecs = jax.tree.map(lambda s: P() if s is None else s, specs)","typeGuard":"def specs_valid(specs) -> bool:\n    return all(s is None or isinstance(s, PartitionSpec) for s in jax.tree.leaves(specs))","tryCatchPattern":"try: shard_map(...) except TypeError as e: if 'was None' in str(e): replace None with P(); else: raise","preventionTips":["Default parameters to P() not None","Normalize config-loaded spec trees with tree.map before calling shard_map"],"tags":["shard-map","partition-spec","none-default","jax","typeerror"],"backgroundTag":"none-instead-of-empty-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}