jax-ml/jax · error · ValueError

{full_name} must form a tree prefix of the corresponding val

Error message

{full_name} must form a tree prefix of the corresponding values (up to pytree node types), but {where} is a tuple while the corresponding part of the values is a leaf; use a single bool there instead

What it means

JAX requires the boolean prefix tree to mirror the structure of the values it annotates. This error fires when the prefix has a tuple at a position where the corresponding value is a leaf (e.g. a single array), so there is nothing for the tuple's children to correspond to.

Source

Thrown at jax/_src/api.py:1900

  return ret

def _saveable_args_flags(saveable_args, treedef) -> list[bool]:
  return tuptree_flags(saveable_args, treedef, 'saveable_args',
                       'the saveable_args argument to jax.vjp')

def _tuptree_flags_rec(prefix, td, name, full_name, path, ret):
  if isinstance(prefix, bool):
    ret.extend([prefix] * td.num_leaves)
    return
  where = name + ''.join(f'[{i}]' for i in path)
  if not isinstance(prefix, tuple):
    raise ValueError(
        f"{full_name} must be a pytree prefix with bool leaves or a "
        f"tuple-tree of bools "
        f"(made of bools and tuples only), but {where} is {prefix!r} of type "
        f"{type(prefix).__name__}")
  if treedef_is_strict_leaf(td):
    raise ValueError(
        f"{full_name} must form a tree prefix of "
        f"the corresponding values (up to pytree node types), but {where} is "
        "a tuple while the corresponding part of the values is a leaf; use "
        "a single bool there instead")
  td_children = td.children()
  if len(prefix) != len(td_children):
    raise ValueError(
        f"{full_name} must form a tree prefix of "
        "the corresponding values (up to pytree node types, so containers "
        f"need only match in their number of children), but {where} has "
        f"{len(prefix)} children while the corresponding container has "
        f"{len(td_children)}")
  for i, (p, td_) in enumerate(zip(prefix, td_children)):
    _tuptree_flags_rec(p, td_, name, full_name, (*path, i), ret)

def _is_ref(x):
  from jax._src.state.types import AbstractRef
  try:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace the tuple at that position with a single bool
2. Check the actual structure of the values with jax.tree_util.tree_structure and mirror it
3. Unwrap or fix the value side if the function returns a leaf where you expected a container

Example fix

// before
prefix = (True, True)   # value here is a single array
// after
prefix = True

import jax.tree_util as jtu
# prefix must not be a tuple where values are a leaf
if isinstance(prefix, tuple) and jtu.tree_structure(values).num_leaves == 1 and not isinstance(values, tuple):
    prefix = prefix[0]

def is_valid_prefix(prefix, values) -> bool:
    if isinstance(prefix, bool): return True
    if isinstance(values, tuple) and isinstance(prefix, tuple):
        return len(prefix) == len(values) and all(is_valid_prefix(p, v) for p, v in zip(prefix, values))
    return False

When it happens

Trigger: Passing (True, True) as a prefix where the value at that position is one array leaf; passing a 1-tuple like (False,) against a scalar output.

Common situations: Mis-counting nesting: values were flattened or wrapped unexpectedly (e.g. a function returning a bare array instead of a pair), so the assumed prefix structure is one level too deep.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {

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