jax-ml/jax · error · TypeError

shard_map in_specs argument must be a pytree of `jax.shardin

Error message

shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but it was None.
Instead of `in_specs=None`, did you mean `in_specs=P()`, where `P = jax.sharding.PartitionSpec`?

What it means

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.

Source

Thrown at jax/_src/shard_map.py:476

          " shard_map is in full manual mode. Got mesh axis names"
          f" {mesh.axis_names}, manual_axes: {manual_axes}, specs: {s}. Please"
          " file a bug at https://github.com/jax-ml/jax/issues.")
    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced):
      raise ValueError(
          f"unreduced in {prefix}_specs {s} can only be used when the mesh"
          " passed to shard_map contains axis names all of type `Explicit`."
          f" Got mesh {mesh}")
    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.reduced):
      raise ValueError(
          f"reduced in {prefix}_specs {s} can only be used when the mesh"
          " passed to shard_map contains axis names all of type `Explicit`."
          f" Got mesh {mesh}")


def _check_specs(error_type: SpecErrorType, specs: Any, manual_axes) -> None:
  from jax._src.hijax import HiPspec
  if error_type == SpecErrorType.input and specs is None:
    raise TypeError(
        "shard_map in_specs argument must be a pytree of "
        "`jax.sharding.PartitionSpec` instances, but it was None.\n"
        "Instead of `in_specs=None`, did you mean `in_specs=P()`, "
        "where `P = jax.sharding.PartitionSpec`?")

  def check_spec(p):
    if isinstance(p, HiPspec):
      return True  # TODO(mattjj,yashkatariya): add user validation method
    if not isinstance(p, PartitionSpec):
      return False
    for names in p.partitions:
      names = (names,) if not isinstance(names, tuple) else names
      for name in names:
        if name is not None and name not in manual_axes:
          return False
    return True

  if all(check_spec(p) for p in tree_leaves(specs)):

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

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

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
from jax.sharding import PartitionSpec as P
specs = jax.tree.map(lambda s: P() if s is None else s, specs)

Type guard

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

Try / catch

try: shard_map(...) except TypeError as e: if 'was None' in str(e): replace None with P(); else: raise

Prevention

Default parameters to P() not None
Normalize config-loaded spec trees with tree.map before calling shard_map

When it happens

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

Common situations: Functions with `in_specs=None` default arguments; dataclass/config fields that default to None and are forwarded into shard_map.

Related errors

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