{"record":{"id":"46902d22b9829606","repo":"jax-ml/jax","slug":"expected-tuple-got-s","errorCode":null,"errorMessage":"Expected tuple, got %s.","messagePattern":"Expected tuple, got (.+?)\\.","errorType":"validation","errorClass":"std::invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1002,"sourceCode":"        --leaf;\n        break;\n\n      case PyTreeKind::kNone:\n        if (!object.is_none()) {\n          throw std::invalid_argument(absl::StrFormat(\n              \"Expected None, got %s.\\n\\n\"\n              \"In previous releases of JAX, flatten-up-to used to \"\n              \"consider None to be a tree-prefix of non-None values. To obtain \"\n              \"the previous behavior, you can usually write:\\n\"\n              \"  jax.tree.map(lambda x, y: None if x is None else f(x, y), a, \"\n              \"b, is_leaf=lambda x: x is None)\",\n              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        break;\n\n      case PyTreeKind::kTuple: {\n        if (!PyTuple_CheckExact(object.ptr())) {\n          throw std::invalid_argument(\n              absl::StrFormat(\"Expected tuple, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        nb::tuple tuple = nb::borrow<nb::tuple>(object);\n        if (tuple.size() != node.arity) {\n          throw std::invalid_argument(absl::StrFormat(\n              \"Tuple arity mismatch: %d != %d; tuple: %s.\", tuple.size(),\n              node.arity, nb::cast<std::string_view>(nb::repr(object))));\n        }\n        for (nb::handle entry : tuple) {\n          agenda.push_back(nb::borrow<nb::object>(entry));\n        }\n        break;\n      }\n\n      case PyTreeKind::kList: {\n        if (!PyList_CheckExact(object.ptr())) {\n          throw std::invalid_argument(","sourceCodeStart":984,"sourceCodeEnd":1020,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L984-L1020","documentation":"JAX/pytree raised this while flattening a Python object against a stored PyTreeDef (treedef) via FlattenUpTo. The treedef records that this node in the tree structure is an exact built-in tuple, but the object being flattened has a different type at that position. It is thrown whenever a function/transform re-applies a structure captured from different data (e.g. jit/pmap with donated or reused arguments, tree_unflatten-like paths, vmap out_axes matching).","triggerScenarios":"Calling an API that flattens a tree up to a prefix treedef (e.g. jax.tree_util.tree_flatten(..., is_leaf=...), FlattenUpTo paths used by jax.jit cache matching / tree_map with mismatched structures) where the treedef node is kTuple (PyTuple_CheckExact) but the runtime object is a list, namedtuple, or other sequence. Subclassing tuple is not enough: the check is exact.","commonSituations":"Passing a list where a tuple was captured at trace time (jit re-trace vs cached path), refactoring data containers from tuple to list (or dataclass) between code revisions, mixing namedtuple and plain tuple, replaying a treedef saved/pickled from an older run against new data.","solutions":["Make the object an exact built-in tuple at that position: convert with tuple(x) before passing it in.","Re-derive the treedef from the current object (re-jit / re-flatten) instead of reusing a treedef captured from earlier data.","If you intentionally changed structure, clear cached compilations (e.g. new jax.jit wrapper) or avoid reusing stored treedefs across structure changes.","Replace tuple nodes with a structure-insensitive container (e.g. dict or custom pytree node registered with jax.tree_util.register_pytree_node) if you need flexible sequence types."],"exampleFix":"# before\nleaves = treedef.flatten_up_to((1, [2, 3]))  # second element was tuple in treedef\n# TypeError/ValueError: Expected tuple, got [2, 3].\n\n# after\nleaves = treedef.flatten_up_to((1, (2, 3)))  # use exact tuple","handlingStrategy":"type-guard","validationCode":"import jax\nexpected = jax.tree_util.tree_structure(reference_tuple)\ndef structure_ok(obj):\n    try:\n        jax.tree_util.tree_flatten(obj)\n        return True\n    except (TypeError, ValueError):\n        return False","typeGuard":"from jax.tree_util import tree_structure\ndef is_exact_tuple_tree(obj, ref):\n    return tree_structure(obj) == tree_structure(ref) and all(\n        isinstance(x, tuple) and type(x) is tuple\n        for x in jax.tree_util.tree_leaves(obj) or [obj]\n    )","tryCatchPattern":"try:\n    treedef.flatten_up_to(obj)\nexcept (ValueError, TypeError) as e:\n    if 'Expected tuple' in str(e):\n        obj = jax.tree_util.tree_map(lambda x: tuple(x) if isinstance(x, list) else x, obj)\n    else:\n        raise","preventionTips":["Keep container types consistent (always tuple, never mix list/tuple) across trace and call sites.","Compare jax.tree_util.tree_structure(obj) == treedef before flattening in debug builds.","Convert at boundaries: tuple(x) for sequence inputs before JAX APIs."],"tags":["jax","pytree","type-mismatch","tuple"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}