{"record":{"id":"4c1ce4dfd3a40de3","repo":"jax-ml/jax","slug":"malformed-pickled-pytreedef-expected-2-tuple","errorCode":null,"errorMessage":"Malformed pickled PyTreeDef, expected 2-tuple","messagePattern":"Malformed pickled PyTreeDef, expected 2-tuple","errorType":"exception","errorClass":"XlaRuntimeError","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1913,"sourceCode":"  treedef.def_static(\n      \"from_node_data_and_children\", &PyTreeDef::FromNodeDataAndChildren,\n      nb::arg(\"registry\"), nb::arg(\"node_data\").none(), nb::arg(\"children\"),\n      \"Reconstructs a pytree from `node_data()` and `children()`.\",\n      nb::sig(\n          // clang-format off\n        \"def from_node_data_and_children(\"\n        \"self, \"\n        \"registry: PyTreeRegistry, \"\n        \"node_data: tuple[type, Any] | None, \"\n        \"children: typing.Iterable[PyTreeDef]\"\n        \") -> PyTreeDef\"\n          // clang-format on\n          ));\n  treedef.def(\"__getstate__\", &PyTreeDef::ToPickle);\n  treedef.def(\"__setstate__\", [](PyTreeDef& t, nb::object o) {\n    nb::tuple pickle = nb::cast<nb::tuple>(o);\n    if (pickle.size() != 2) {\n      throw xla::XlaRuntimeError(\n          \"Malformed pickled PyTreeDef, expected 2-tuple\");\n    }\n    auto registry = nb::cast<nb_class_ptr<PyTreeRegistry>>(pickle[0]);\n    new (&t) PyTreeDef(registry);\n    t.FromPickle(pickle[1]);\n  });\n\n  nb::class_<SequenceKey> sequence_key(\n      pytree, \"SequenceKey\", nb::sig(\"class SequenceKey(typing.Hashable)\"));\n  sequence_key.def(nb::init<int>(), nb::arg(\"idx\"));\n  sequence_key.def(\"__str__\", &SequenceKey::ToString);\n  sequence_key.def(\"__repr__\", &SequenceKey::ToReprString);\n  sequence_key.def(\"__eq__\", &SequenceKey::Equals);\n  sequence_key.def(\"__hash__\", [](const SequenceKey& key) {\n    return key.idx() + kSequenceKeyHashSalt;\n  });\n  sequence_key.def_prop_ro(\"idx\", &SequenceKey::idx);\n  sequence_key.def_prop_ro_static(\"__match_args__\", &SequenceKey::MatchArgs);","sourceCodeStart":1895,"sourceCodeEnd":1931,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1895-L1931","documentation":"PyTreeDef.__setstate__ expects the pickled state to be a 2-tuple of (registry, treedef_data). If the pickle contains anything other than exactly 2 elements, this XlaRuntimeError is thrown.","triggerScenarios":"Unpickling a PyTreeDef whose __getstate__ output was tampered with, produced by a different jaxlib version with a different pickle layout, or manually constructed.","commonSituations":"Loading old pickles saved with an incompatible JAX version; hand-editing or wrapping pickled state; corrupted pickle files.","solutions":["Re-create the pickle with the current jaxlib version instead of loading stale ones","Verify the pickle contains a genuine PyTreeDef, not a modified tuple","Match JAX versions between save and load environments"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"import pickle\ndef valid_pytreedef_state(o) -> bool:\n    return isinstance(o, tuple) and len(o) == 2","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never hand-edit pickled PyTreeDef state","Regenerate pickles after jaxlib upgrades"],"tags":["jax","pytree","pickle"],"backgroundTag":"pickle-format-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}