jax-ml/jax · error · TypeError

shard_map {prefix}_specs argument must be a pytree of `jax.s

Error message

shard_map {prefix}_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but:

{msgs}

Check the {prefix}_specs values passed to shard_map.

What it means

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).

Source

Thrown at jax/_src/shard_map.py:511

  if all(check_spec(p) for p in tree_leaves(specs)):
    return
  prefix = 'in' if error_type == SpecErrorType.input else 'out'
  msgs = [f"  {prefix}_specs{keystr(key)} is {x} of type {type(x).__name__}, "
          for key, x in generate_key_paths(specs) if not isinstance(x, P)]
  if not msgs:
    for key, p in generate_key_paths(specs):
      for names in p:
        names = (names,) if not isinstance(names, tuple) else names
        for name in names:
          if name is not None and name not in manual_axes:
            msgs.append(f"  {prefix}_specs{keystr(key)} refers to {repr(name)}")
    raise ValueError(
        f"shard_map {prefix}_specs argument must refer to an axis "
        f"marked as manual ({manual_axes}), but:\n\n"
        + '\n\n'.join(msgs) + '\n\n'
        f"Check the {prefix}_specs values passed to shard_map.")
  raise TypeError(
      f"shard_map {prefix}_specs argument must be a pytree of "
      f"`jax.sharding.PartitionSpec` instances, but:\n\n"
      + '\n\n'.join(msgs) + '\n\n'
      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)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap every spec leaf in PartitionSpec (P), e.g. P('data','model') instead of ('data','model')
2. Sanitize with jax.tree.map(lambda s: s if isinstance(s, P) else P(*s), specs)
3. Validate spec trees before calling shard_map

Example fix

// before
shard_map(f, mesh, xs, in_specs=(('data',), None))
// after
from jax.sharding import PartitionSpec as P
shard_map(f, mesh, xs, in_specs=(P('data'), P()))

from jax.sharding import PartitionSpec as P
specs = jax.tree.map(lambda s: s if isinstance(s, P) else P(*s) if isinstance(s, (tuple, list)) else s, specs)
assert all(isinstance(s, P) for s in jax.tree.leaves(specs))

def is_spec_pytree(specs) -> bool:
    return all(isinstance(s, PartitionSpec) for s in jax.tree.leaves(specs))

try: shard_map(...) except TypeError as e: if 'pytree of' in str(e): coerce leaves to PartitionSpec and retry; else: raise

When it happens

Trigger: Passing raw tuples like in_specs=(('data','model'),) instead of PartitionSpec objects, or mixing strings/other objects into the spec pytree.

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

Related errors

• shard_map in_specs argument must be a pytree of `jax.shardin
• 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

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