jax-ml/jax · error · ValueError

pmap requires at least one argument with a mapped axis.

Error message

pmap requires at least one argument with a mapped axis.

What it means

pmap infers the mapped axis size from arguments whose in_axes entry is not None. If every argument has in_axes=None (fully broadcast) and no explicit axis_size is available, there is no way to determine how many devices to map over.

Source

Thrown at jax/_src/pmap.py:498

    args: Flat list of arguments.
    in_axes: Flat tuple of axis indices (int or None for each arg).

  Returns:
    The size of the mapped axis.

  Raises:
    ValueError: If no args have a mapped axis.
  """
  if args and in_axes:
    # Fast path: check first arg/axis (most common case).
    if in_axes[0] is not None and hasattr(args[0], "shape"):
      return int(args[0].shape[in_axes[0]])
    # Slow path: scan for first mapped arg.
    if isinstance(in_axes, tuple):
      for arg, ax in zip(args, in_axes):
        if ax is not None and hasattr(arg, "shape"):
          return int(arg.shape[ax])
  raise ValueError("pmap requires at least one argument with a mapped axis.")


def _pmap_wrap_init(f, static_broadcasted_tuple):
  """Create a wrapped function with DebugInfo for pmap.

  Args:
    f: The function to wrap.
    static_broadcasted_tuple: Tuple of static argument indices.

  Returns:
    A lu.WrappedFun ready for pmap.
  """
  # Compute arg_names from signature, excluding static argnums
  if (signature := fun_signature(f)) is not None:
    static_set = frozenset(static_broadcasted_tuple)
    arg_names = tuple(
        name
        for i, name in enumerate(signature.parameters.keys())

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set in_axes to an int (e.g. 0) for at least one array argument
2. Pass `axis_size=N` to pmap explicitly so the mapped size is known
3. Restructure to jit + sharding if no argument actually varies per device

Example fix

# before
f = jax.pmap(fn, in_axes=None)
f(x)
# after
f = jax.pmap(fn, in_axes=0)
f(x)

from jax.tree_util import tree_leaves
has_mapped = any(a is not None for a in tree_leaves(in_axes))
assert has_mapped or axis_size is not None, 'pmap needs a mapped axis or axis_size'

When it happens

Trigger: `jax.pmap(f, in_axes=None)(x)` where all leaves of in_axes are None, and no axis_size given; also when all args lack a `.shape` attribute (e.g. scalars/static objects).

Common situations: Broadcasting a constant computation over devices; passing static/metadata-only arguments; wrong in_axes tuple ordering so mapped axes align with non-array args.

Related errors

• {name} must have at least one non-None value in in_axes or a
• pmap in_axes must be an int, None, or (nested) container wit
• pmap in_axes: value not a valid prefix (dynamic message incl
• Specified axis_size {axis_size} doesn't match received axis_
• {name} wrapped function must be passed at least one argument

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