jax-ml/jax · error · ValueError

shard_map in_specs rank error (msg from _spec_rank_error)

Error message

shard_map in_specs rank error (msg from _spec_rank_error)

What it means

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.

Source

Thrown at jax/_src/shard_map.py:533

      f"Check the {prefix}_specs values passed to shard_map.")

class NoFail:
  def __repr__(self):
    return "NoFail()"

no_fail = NoFail()

def _check_specs_vs_args(
    f: Callable, mesh: Mesh | AbstractMesh, in_tree: PyTreeDef, in_specs: Specs,
    dyn_argnums: Sequence[int], in_specs_flat: Sequence[P],
    xs: Sequence) -> None:
  in_avals = map(core.shaped_abstractify, xs)
  fail = [a if isinstance(p, P) and len(p) > a.ndim else no_fail
          for p, a in zip(in_specs_flat, in_avals)]
  if any(f is not no_fail for f in fail):
    fail = _expand_fail(in_tree, dyn_argnums, fail)
    msg = _spec_rank_error(SpecErrorType.input, f, in_tree, in_specs, fail)
    raise ValueError(msg)
  bad = lambda a, d, ns: a.shape[d] % prod(mesh.shape[n] for n in ns)
  fail = [a if (isinstance(s, P) and
                any(bad(a, d, ns) for d, ns in _spec_to_names(s).items()))
          else no_fail for a, s in zip(in_avals, in_specs_flat)]
  if any(f is not no_fail for f in fail):
    fail = _expand_fail(in_tree, dyn_argnums, fail)
    msg = _spec_divisibility_error(f, mesh, in_tree, in_specs, fail)
    raise ValueError(msg)

def _expand_fail(in_tree: PyTreeDef, dyn_argnums: Sequence[int],
                 fail: Sequence[core.ShapedArray | NoFail]
                 ) -> list[core.ShapedArray | NoFail]:
  fail_: list[core.ShapedArray | NoFail] = [no_fail] * in_tree.num_leaves
  for i, f in zip(dyn_argnums, fail):
    fail_[i] = f
  return fail_

def _spec_rank_error(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match spec rank to each array's ndim (use P() or None to replicate lower-rank arrays)
2. Build per-leaf specs from actual shapes: tree_map over avals
3. Pad singleton dims or reshape the array if the spec is intentional

Example fix

// before
shard_map(f, mesh, (w, b), in_specs=(P('i','j'), P('i','j')))  # b is 1D
// after
shard_map(f, mesh, (w, b), in_specs=(P('i','j'), P()))

def rank_ok(tree, args, specs):
    avals = map(jax.core.shaped_abstractify, jax.tree.leaves(args))
    return all(len(p) <= a.ndim for p, a in zip(jax.tree.leaves(specs), avals))

def spec_fits(spec, ndim) -> bool:
    return isinstance(spec, P) and len(spec) <= ndim

try: shard_map(...) except ValueError as e: if 'rank' in str(e): demote offending specs to P(); else: raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• in_specs passed to shard_map: {s} does not match the specs o
• Invalid spec: {spec}
• unreduced in {prefix}_specs {s} can only be used when the me
• reduced in {prefix}_specs {s} can only be used when the mesh
• shard_map in_specs argument must be a pytree of `jax.shardin

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/38ca50b12030c71b. Report an issue: GitHub.