{"record":{"id":"0bc9cc94e0e8bdd8","repo":"jax-ml/jax","slug":"full-name-must-be-a-pytree-prefix-with-bool-leav","errorCode":null,"errorMessage":"{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__}","messagePattern":"(.+?) must be a pytree prefix with bool leaves or a tuple-tree of bools \\(made of bools and tuples only\\), but (.+?) is (.+?) of type (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1894,"sourceCode":"    pass\n  else:\n    if all(isinstance(f, bool) for f in flags):\n      return list(flags)\n  ret: list[bool] = []\n  _tuptree_flags_rec(prefix, treedef, name, full_name, (), ret)\n  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)}\")","sourceCodeStart":1876,"sourceCodeEnd":1912,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1876-L1912","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Replace any non-bool, non-tuple leaves with bools (0/1 -> False/True)","Convert lists to tuples: [True, False] -> (True, False)","Broadcast a single bool instead of nesting: (True, True) -> True"],"exampleFix":"// before\nflags = [True, False]\n// after\nflags = (True, False)","handlingStrategy":"validation","validationCode":"def check_bool_tuptree(p):\n    if isinstance(p, bool): return True\n    if isinstance(p, tuple):\n        return all(check_bool_tuptree(x) for x in p)\n    return False\nassert check_bool_tuptree(prefix)","typeGuard":"def is_bool_tuptree(p) -> bool:\n    return (isinstance(p, bool) or\n            (isinstance(p, tuple) and all(is_bool_tuptree(x) for x in p)))","tryCatchPattern":null,"preventionTips":["Always use tuples, never lists, for prefix trees","Use plain Python True/False, not 0/1 or numpy bools"],"tags":["jax","pytree","validation","bool-prefix"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}