{"record":{"id":"80c0a2b1a7981403","repo":"jax-ml/jax","slug":"full-name-must-form-a-tree-prefix-of-the-corresp","errorCode":null,"errorMessage":"{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","messagePattern":"(.+?) must form a tree prefix of the corresponding values \\(up to pytree node types\\), but (.+?) is a tuple while the corresponding part of the values is a leaf; use a single bool there instead","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1900,"sourceCode":"  return ret\n\ndef _saveable_args_flags(saveable_args, treedef) -> list[bool]:\n  return tuptree_flags(saveable_args, treedef, 'saveable_args',\n                       'the saveable_args argument to jax.vjp')\n\ndef _tuptree_flags_rec(prefix, td, name, full_name, path, ret):\n  if isinstance(prefix, bool):\n    ret.extend([prefix] * td.num_leaves)\n    return\n  where = name + ''.join(f'[{i}]' for i in path)\n  if not isinstance(prefix, tuple):\n    raise ValueError(\n        f\"{full_name} must be a pytree prefix with bool leaves or a \"\n        f\"tuple-tree of bools \"\n        f\"(made of bools and tuples only), but {where} is {prefix!r} of type \"\n        f\"{type(prefix).__name__}\")\n  if treedef_is_strict_leaf(td):\n    raise ValueError(\n        f\"{full_name} must form a tree prefix of \"\n        f\"the corresponding values (up to pytree node types), but {where} is \"\n        \"a tuple while the corresponding part of the values is a leaf; use \"\n        \"a single bool there instead\")\n  td_children = td.children()\n  if len(prefix) != len(td_children):\n    raise ValueError(\n        f\"{full_name} must form a tree prefix of \"\n        \"the corresponding values (up to pytree node types, so containers \"\n        f\"need only match in their number of children), but {where} has \"\n        f\"{len(prefix)} children while the corresponding container has \"\n        f\"{len(td_children)}\")\n  for i, (p, td_) in enumerate(zip(prefix, td_children)):\n    _tuptree_flags_rec(p, td_, name, full_name, (*path, i), ret)\n\ndef _is_ref(x):\n  from jax._src.state.types import AbstractRef\n  try:","sourceCodeStart":1882,"sourceCodeEnd":1918,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1882-L1918","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Replace the tuple at that position with a single bool","Check the actual structure of the values with jax.tree_util.tree_structure and mirror it","Unwrap or fix the value side if the function returns a leaf where you expected a container"],"exampleFix":"// before\nprefix = (True, True)   # value here is a single array\n// after\nprefix = True","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\n# prefix must not be a tuple where values are a leaf\nif isinstance(prefix, tuple) and jtu.tree_structure(values).num_leaves == 1 and not isinstance(values, tuple):\n    prefix = prefix[0]","typeGuard":"def is_valid_prefix(prefix, values) -> bool:\n    if isinstance(prefix, bool): return True\n    if isinstance(values, tuple) and isinstance(prefix, tuple):\n        return len(prefix) == len(values) and all(is_valid_prefix(p, v) for p, v in zip(prefix, values))\n    return False","tryCatchPattern":null,"preventionTips":["Print jax.tree_util.tree_structure(values) when unsure","Prefer broadcasting a single bool where possible"],"tags":["jax","pytree","tree-prefix"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}