{"record":{"id":"2e0c34aacd53b986","repo":"jax-ml/jax","slug":"pytree-serialization-too-large-to-deserialize","errorCode":null,"errorMessage":"Pytree serialization too large to deserialize.","messagePattern":"Pytree serialization too large to deserialize\\.","errorType":"exception","errorClass":"XlaRuntimeError","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1883,"sourceCode":"  treedef.def(\"__ne__\", [](const PyTreeDef& a, nb::object b) {\n    return nb::isinstance<PyTreeDef>(b) && a != nb::cast<PyTreeDef>(b);\n  });\n  treedef.def(\"__hash__\", [](const PyTreeDef& t) -> Py_hash_t {\n    return AbslHashToPythonHash(absl::HashOf(t));\n  });\n  treedef.def(\"serialize_using_proto\", [](const PyTreeDef& a) {\n    PyTreeDefProto result;\n    a.SerializeTo(result);\n    std::string serialized = result.SerializeAsString();\n    return nb::bytes(serialized.data(), serialized.size());\n  });\n  treedef.def_static(\n      \"deserialize_using_proto\",\n      [](nb_class_ptr<PyTreeRegistry> registry, nb::bytes data) {\n        PyTreeDefProto input;\n        std::string_view serialized(data.c_str(), data.size());\n        if (serialized.size() > std::numeric_limits<int>::max()) {\n          throw xla::XlaRuntimeError(\n              \"Pytree serialization too large to deserialize.\");\n        }\n        if (!input.ParseFromArray(serialized.data(), serialized.size())) {\n          throw xla::XlaRuntimeError(\"Could not deserialize PyTreeDefProto.\");\n        }\n        return PyTreeDef::DeserializeFrom(std::move(registry), input);\n      },\n      nb::arg(\"registry\"), nb::arg(\"data\"));\n  treedef.def(\"node_data\", &PyTreeDef::GetNodeData,\n              \"Returns None if a leaf-pytree, else (type, node_data)\",\n              nb::sig(\"def node_data(self) -> tuple[type, Any] | None\"));\n  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(\"","sourceCodeStart":1865,"sourceCodeEnd":1901,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1865-L1901","documentation":"The static binding deserialize_using_proto rejects payloads larger than INT_MAX bytes because protobuf's ParseFromArray takes an int. Any serialized pytree proto above ~2GB cannot be deserialized.","triggerScenarios":"Calling deserialize_using_proto with a bytes object whose length exceeds std::numeric_limits<int>::max() (2147483647 bytes) on a 64-bit build where size_t is wider.","commonSituations":"Extremely large/deep pytrees (millions of nodes) serialized and shipped between processes; rarely hit except with machine-generated nested structures or accidentally re-serialized (nested) payloads.","solutions":["Reduce the pytree size: flatten into leaves + treedef separately, or chunk the structure","Check for accidental double serialization (bytes of bytes) inflating size","Avoid serializing gigantic static structures as pytrees; use arrays"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"import sys\nMAX = 2**31 - 1\ndef safe_deserialize(registry, data: bytes):\n    if len(data) > MAX:\n        raise ValueError(f'payload {len(data)} exceeds protobuf int limit')\n    return pytree.PyTreeDef.deserialize_using_proto(registry, data)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep serialized pytrees small; store large data as arrays, not structure","Watch for accidental double serialization"],"tags":["jax","pytree","protobuf","size-limit"],"backgroundTag":"payload-too-large","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}