jax-ml/jax · error · ValueError

pmap in_axes: value not a valid prefix (dynamic message incl

Error message

pmap in_axes: value not a valid prefix (dynamic message including 'The full pytree here is the tuple of arguments passed positionally to the pmapped function...' and 'Check that the value of the `in_axes` argument to `pmap` is a tree prefix of the tuple of arguments passed positionally to the pmapped function.')

What it means

pmap flattens `in_axes` against the flattened tree of positional arguments; if in_axes is not a valid pytree prefix of the args tuple (wrong arity, wrong nesting, or extra levels), flattening fails and pmap raises a detailed prefix-mismatch ValueError.

Source

Thrown at jax/_src/pmap.py:645

          "\n\nThe 'full pytree' here is the tuple of arguments passed "
          "positionally to the pmapped function, and the value of `in_axes` "
          "must be a tree prefix of that tuple. But it was not a prefix."
      )
      if kwargs:
        msg += (
            "\n\nWhen some arguments are passed by keyword to the pmapped "
            "function, they are not included in the comparison to `in_axes`. "
            "Instead, each argument passed by keyword is mapped over its "
            "leading axis. See the description of `in_axes` in the `pmap` "
            "docstring: "
            "https://docs.jax.dev/en/latest/_autosummary/jax.pmap.html#jax.pmap"
        )
      msg += (
          "\n\nCheck that the value of the `in_axes` argument to `pmap` "
          "is a tree prefix of the tuple of arguments passed positionally to "
          "the pmapped function."
      )
      raise ValueError(msg) from None

  return in_axes_flat


def _get_donated_invars(donate_tuple, in_tree, num_flat_args):
  """Compute donation vector for arguments.

  Args:
    donate_tuple: Tuple of donated argument indices.
    in_tree: PyTreeDef of input structure.
    num_flat_args: Number of flat arguments.

  Returns:
    Tuple of bools indicating which flat args are donated.
  """

  if donate_tuple and not config.debug_nans.value:
    return donation_vector(donate_tuple, (), in_tree)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make in_axes exactly mirror the structure of the positional arguments tuple (or use a scalar like 0)
2. Pass all mapped arguments positionally
3. Use `jax.tree_util.tree_structure(args)` to debug and match structures

Example fix

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

from jax.tree_util import tree_flatten
ax_leaves = tree_flatten(in_axes)[0]
arg_leaves = tree_flatten(args)[0]
# in_axes must be a prefix: flatten_axes would fail; quick arity sanity check
assert len(ax_leaves) <= len(arg_leaves), 'in_axes arity exceeds args'

When it happens

Trigger: Calling a pmapped function with a different number/order of arguments than the in_axes tree describes, e.g. `in_axes=(0,0)` but calling with 3 args, or nested pytrees in args not matched by in_axes structure.

Common situations: Optional args passed sometimes and not others; refactored function signatures; passing kwargs (which are not part of the positional tuple) while in_axes assumed them positional.

Related errors

• pmap in_axes must be an int, None, or (nested) container wit
• pmap requires at least one argument with a mapped axis.
• {name} was requested to map a value of non-array type {core.
• {name} was requested to map its argument along axis {axis},
• {name} must have at least one non-None value in in_axes or a

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