{"record":{"id":"38ca50b12030c71b","repo":"jax-ml/jax","slug":"shard-map-in-specs-rank-error-msg-from-spec-rank","errorCode":null,"errorMessage":"shard_map in_specs rank error (msg from _spec_rank_error)","messagePattern":"shard_map in_specs rank error \\(msg from _spec_rank_error\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":533,"sourceCode":"      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)]\n  if any(f is not no_fail for f in fail):\n    fail = _expand_fail(in_tree, dyn_argnums, fail)\n    msg = _spec_rank_error(SpecErrorType.input, f, in_tree, in_specs, fail)\n    raise ValueError(msg)\n  bad = lambda a, d, ns: a.shape[d] % prod(mesh.shape[n] for n in ns)\n  fail = [a if (isinstance(s, P) and\n                any(bad(a, d, ns) for d, ns in _spec_to_names(s).items()))\n          else no_fail for a, s in zip(in_avals, in_specs_flat)]\n  if any(f is not no_fail for f in fail):\n    fail = _expand_fail(in_tree, dyn_argnums, fail)\n    msg = _spec_divisibility_error(f, mesh, in_tree, in_specs, fail)\n    raise ValueError(msg)\n\ndef _expand_fail(in_tree: PyTreeDef, dyn_argnums: Sequence[int],\n                 fail: Sequence[core.ShapedArray | NoFail]\n                 ) -> list[core.ShapedArray | NoFail]:\n  fail_: list[core.ShapedArray | NoFail] = [no_fail] * in_tree.num_leaves\n  for i, f in zip(dyn_argnums, fail):\n    fail_[i] = f\n  return fail_\n\ndef _spec_rank_error(","sourceCodeStart":515,"sourceCodeEnd":551,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L515-L551","documentation":"Each PartitionSpec entry consumes one dimension of the corresponding array; if a spec has more entries than the array has dimensions (spec rank > aval rank), the sharding is ill-defined and shard_map raises this ValueError with per-argument details.","triggerScenarios":"shard_map(f, mesh, x) where x has shape (8,) but in_specs=P('a','b') (rank 2 > 1); commonly from 1D bias/offset arrays in a pytree sharing a spec meant for 2D weights.","commonSituations":"Applying a single spec tree to a heterogeneous parameter pytree; adding a batch dimension on one side but not the other; off-by-one dimension counts after refactoring.","solutions":["Match spec rank to each array's ndim (use P() or None to replicate lower-rank arrays)","Build per-leaf specs from actual shapes: tree_map over avals","Pad singleton dims or reshape the array if the spec is intentional"],"exampleFix":"// before\nshard_map(f, mesh, (w, b), in_specs=(P('i','j'), P('i','j')))  # b is 1D\n// after\nshard_map(f, mesh, (w, b), in_specs=(P('i','j'), P()))","handlingStrategy":"validation","validationCode":"def rank_ok(tree, args, specs):\n    avals = map(jax.core.shaped_abstractify, jax.tree.leaves(args))\n    return all(len(p) <= a.ndim for p, a in zip(jax.tree.leaves(specs), avals))","typeGuard":"def spec_fits(spec, ndim) -> bool:\n    return isinstance(spec, P) and len(spec) <= ndim","tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'rank' in str(e): demote offending specs to P(); else: raise","preventionTips":["Generate specs per-leaf from actual array shapes with tree.map","Unit-test spec trees against parameter pytrees"],"tags":["shard-map","rank-mismatch","partition-spec","jax"],"backgroundTag":"shape-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}