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, so containers need only match in their number of children), but {where} has {len(prefix)} children while the corresponding container has {len(td_children)}

What it means

The boolean prefix tuple must have the same number of children as the corresponding container in the values. JAX raises this when the tuple lengths diverge, meaning the prefix is not a valid tree prefix even though node kinds match.

Source

Thrown at jax/_src/api.py:1907

  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:
    return isinstance(typeof(x), AbstractRef)
  except:
    return False

def _is_ref_aval(a):
  from jax._src.state.types import AbstractRef
  return isinstance(a, AbstractRef)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the values structure with jax.tree_util.tree_structure(values) and count children
2. Update the prefix tuple length to match the container's length
3. Use a single bool to broadcast over all children of that container

Example fix

// before
prefix = (True,)          # values = (a, b)
// after
prefix = True             # or (True, False)

def prefix_len_ok(prefix, values):
    if isinstance(prefix, bool): return True
    return isinstance(values, tuple) and len(prefix) == len(values)

When it happens

Trigger: Passing (True,) against a values tuple of length 2, or (True, False, True) against a pair; any arity mismatch between prefix tuple and container.

Common situations: Adding/removing a return value from a function without updating the prefix; assuming a container has a different length (e.g. treating (a, b) as a 3-tuple).

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/37b2ef5503c3215f. Report an issue: GitHub.