{"record":{"id":"7e2beb94c3f67414","repo":"jax-ml/jax","slug":"list-arity-mismatch-d-d-list-s","errorCode":null,"errorMessage":"List arity mismatch: %d != %d; list: %s.","messagePattern":"List arity mismatch: (.+?) != (.+?); list: (.+?)\\.","errorType":"validation","errorClass":"std::invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1026,"sourceCode":"          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(\n              \"List arity mismatch: %d != %d; list: %s.\", list.size(),\n              node.arity, nb::cast<std::string_view>(nb::repr(object))));\n        }\n        for (nb::handle entry : list) {\n          agenda.push_back(nb::borrow<nb::object>(entry));\n        }\n        break;\n      }\n\n      case PyTreeKind::kDict: {\n        if (!PyDict_CheckExact(object.ptr())) {\n          throw std::invalid_argument(\n              absl::StrFormat(\"Expected dict, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        nb::dict dict = nb::borrow<nb::dict>(object);\n        std::vector<nb::object> keys = GetSortedPyDictKeys(dict.ptr());\n        if (!IsSortedPyDictKeysEqual(keys, node.sorted_dict_keys)) {","sourceCodeStart":1008,"sourceCodeEnd":1044,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1008-L1044","documentation":"PyTreeDef::FlattenUpTo in JAX detected a list node whose length differs from the arity stored in the treedef. Exact-length matching is required because the treedef's prefix must structurally match the object being flattened.","triggerScenarios":"Flattening an object against a prefix treedef whose list node has arity N while the runtime list has a different number of elements; typical in jit-compiled function reuse, tree_map over trees where one side's list length changed, or replaying stored treedefs on batch data of different size.","commonSituations":"Variable batch sizes or sequence lengths stored in lists and passed to a jitted function whose cache expects the traced length, dynamic list building (append in loops) producing inconsistent lengths, editing datasets/config lists between runs.","solutions":["Pad/truncate the list to the arity reported in the message (object size vs node arity).","Wrap the jitted function so lists of varying length trigger re-trace appropriately (e.g. pass length as a jnp scalar arg, or convert to jnp arrays with batch dim) instead of relying on cached treedefs.","If lengths vary by design, recompute the treedef per call or register a custom pytree node that treats length as dynamic aux data.","Add assertions on len() before calling JAX APIs to fail early with a clear message."],"exampleFix":"# before\nbatch = [x1, x2]                  # treedef expects 3 elements\n# after\nbatch = [x1, x2, x3]              # or pad: batch += [dummy]*(3-len(batch))","handlingStrategy":"validation","validationCode":"expected_len = 3\nif len(batch_list) != expected_len:\n    batch_list = (batch_list + [pad] * expected_len)[:expected_len]","typeGuard":"def list_arity_ok(lst: list, arity: int) -> bool:\n    return type(lst) is list and len(lst) == arity","tryCatchPattern":"try:\n    treedef.flatten_up_to(obj)\nexcept ValueError as e:\n    if 'List arity mismatch' in str(e):\n        # resize then retry once\n        raise\n    raise","preventionTips":["Assert len() on variable-length lists before jitted calls.","Design variable-length data as jnp arrays with a batch dimension instead of Python lists.","Re-jit or pass length as an argument when sequence length changes."],"tags":["jax","pytree","arity-mismatch","list"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}