{"record":{"id":"387287f356c682e4","repo":"jax-ml/jax","slug":"tuple-arity-mismatch-d-d-tuple-s","errorCode":null,"errorMessage":"Tuple arity mismatch: %d != %d; tuple: %s.","messagePattern":"Tuple arity mismatch: (.+?) != (.+?); tuple: (.+?)\\.","errorType":"validation","errorClass":"std::invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1008,"sourceCode":"              \"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(\n              absl::StrFormat(\"Expected list, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        nb::list list = nb::borrow<nb::list>(object);\n        if (list.size() != node.arity) {\n          throw std::invalid_argument(absl::StrFormat(","sourceCodeStart":990,"sourceCodeEnd":1026,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L990-L1026","documentation":"During PyTreeDef::FlattenUpTo in JAX's pytree C++ extension, the object at this position is an exact tuple but its length does not equal the arity recorded in the treedef node. JAX enforces strict structural equality of prefixes, so a tuple of the wrong length is rejected even if elements are compatible.","triggerScenarios":"Flattening an object against a prefix treedef whose tuple node has arity N while the runtime tuple has M != N elements; e.g. jax.jit argument donated/reused mismatch, tree_unflatten-then-modify, calling FlattenUpTo-based internal paths (jit cache hits, vmap in_axes/out_axes traversal) with a shortened/lengthened tuple.","commonSituations":"Appending/removing an element from a config tuple between runs while reusing a compiled/cached function, off-by-one tuple construction (trailing comma adding a 1-tuple), conditional code paths building tuples of different lengths, pickled treedefs from older data shape.","solutions":["Fix the tuple length at that position to match the treedef arity shown in the message (%d != %d; first is the object's size).","If the data legitimately changed size, do not reuse the old treedef: re-trace/re-jit or recompute the treedef from the new object.","Build tuples deterministically (e.g. tuple(fixed_length_iterable)) instead of conditional appends.","For variable-length data, use a registered custom pytree node that stores length in aux_data or use lists/dicts with matching structure."],"exampleFix":"# before\ntreedef = jax.tree_util.tree_structure(((0, 0), (0, 0)))\nobj = ((1, 2), (3,))          # inner tuple arity 1 != 2\n# after\nobj = ((1, 2), (3, 4))        # arity matches","handlingStrategy":"validation","validationCode":"node_arity = 2  # from treedef / error message\nassert len(my_tuple) == node_arity, f'expected {node_arity}, got {len(my_tuple)}'","typeGuard":"def tuple_matches(t: tuple[int, ...], arity: int) -> bool:\n    return type(t) is tuple and len(t) == arity","tryCatchPattern":"try:\n    treedef.flatten_up_to(obj)\nexcept ValueError as e:\n    if 'Tuple arity mismatch' in str(e):\n        raise ValueError(f'reshape input to arity in message: {e}') from e\n    raise","preventionTips":["Construct tuples with fixed-length comprehensions: tuple(x for _ in range(n)).","Pass sizes explicitly instead of encoding them in tuple length; re-trace when data shape changes.","Add unit tests asserting tree_structure equality between producer and consumer."],"tags":["jax","pytree","arity-mismatch","tuple"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}