jax-ml/jax · error · ValueError

{full_name} must be a pytree prefix with bool leaves or a tu

Error message

{full_name} must be a pytree prefix with bool leaves or a tuple-tree of bools (made of bools and tuples only), but {where} is {prefix!r} of type {type(prefix).__name__}

What it means

Raised by JAX when a boolean prefix argument (like has_aux-style flags, e.g. in jax.checkpoint or similar prefix APIs) contains a leaf that is neither a bool nor a tuple. JAX validates that the prefix tree is built strictly out of booleans and tuples before broadcasting it against the flattened values.

Source

Thrown at jax/_src/api.py:1894

    pass
  else:
    if all(isinstance(f, bool) for f in flags):
      return list(flags)
  ret: list[bool] = []
  _tuptree_flags_rec(prefix, treedef, name, full_name, (), ret)
  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)}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace any non-bool, non-tuple leaves with bools (0/1 -> False/True)
2. Convert lists to tuples: [True, False] -> (True, False)
3. Broadcast a single bool instead of nesting: (True, True) -> True

Example fix

// before
flags = [True, False]
// after
flags = (True, False)

def check_bool_tuptree(p):
    if isinstance(p, bool): return True
    if isinstance(p, tuple):
        return all(check_bool_tuptree(x) for x in p)
    return False
assert check_bool_tuptree(prefix)

def is_bool_tuptree(p) -> bool:
    return (isinstance(p, bool) or
            (isinstance(p, tuple) and all(is_bool_tuptree(x) for x in p)))

When it happens

Trigger: Passing e.g. in_axes='latents' (a string), a list [True, False], or an int where a tuple-tree of bools is expected (e.g. a prefix argument to vjp/jvp-related or checkpoint APIs that call tuptree_flags).

Common situations: Using a list instead of a tuple for nested prefixes; passing a string axis spec where only bools/tuples are allowed; passing None or 0/1 instead of True/False.

Related errors

• smap in_axes must be an int, None, jax.sharding.Infer, or (n
• 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

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