{"record":{"id":"75028fbfa7a329a7","repo":"jax-ml/jax","slug":"malformed-pytree-proto-invalid-node-type","errorCode":null,"errorMessage":"Malformed pytree proto (invalid node type)","messagePattern":"Malformed pytree proto \\(invalid node type\\)","errorType":"validation","errorClass":"invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1561,"sourceCode":"        break;\n      case PyTreeNodeType::PY_TREE_KIND_NONE:\n        node.kind = PyTreeKind::kNone;\n        break;\n      case PyTreeNodeType::PY_TREE_KIND_TUPLE:\n        node.kind = PyTreeKind::kTuple;\n        break;\n      case PyTreeNodeType::PY_TREE_KIND_DICT:\n        node.kind = PyTreeKind::kDict;\n        for (uint32_t str_id : node_proto.dict_keys().str_id()) {\n          if (str_id >= interned_strings.size()) {\n            throw std::invalid_argument(\n                \"Malformed pytree proto (dict_key out of range).\");\n          }\n          node.sorted_dict_keys.push_back(interned_strings.at(str_id));\n        }\n        break;\n      default:\n        throw std::invalid_argument(\n            \"Malformed pytree proto (invalid node type)\");\n        break;\n    }\n  }\n  result->SetNumLeavesAndNumNodes();\n  return result;\n}\n\nstd::optional<std::pair<nb::object, nb::object>> PyTreeDef::GetNodeData()\n    const {\n  if (traversal_.empty()) {\n    throw std::logic_error(\"empty PyTreeDef traversal.\");\n  }\n  auto builtin_type = [](PyTypeObject* type_obj) {\n    return nb::borrow<nb::object>(reinterpret_cast<PyObject*>(type_obj));\n  };\n  const auto& node = traversal_.back();\n  switch (node.kind) {","sourceCodeStart":1543,"sourceCodeEnd":1579,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1543-L1579","documentation":"Thrown by PyTreeDef::DeserializeFrom when a serialized PyTreeDefProto contains a node whose 'type' oneof field doesn't match any known PyTreeKind (e.g. not one of the defined node_type cases). It means the serialized pytree structure is corrupt or was produced by an incompatible jaxlib version.","triggerScenarios":"Calling jaxlib.pytree.PyTreeDef.deserialize_using_proto (or jax.tree_util APIs that rebuild treedefs from protos) with bytes that fail to parse into a valid node type, e.g. hand-crafted, truncated, or version-mismatched proto data.","commonSituations":"Passing arbitrary bytes instead of output from serialize_using_proto; deserializing pytrees serialized by a newer/older JAX version whose proto schema added node kinds; corrupted payloads over the wire or in storage.","solutions":["Verify the payload was produced by the matching jaxlib version's serialize_using_proto on both ends","Upgrade/downgrade both producer and consumer to the same jaxlib/jax version","Round-trip test: serialize a simple pytree and deserialize it to validate the pipeline","If persisting pytrees, re-export them after upgrading JAX"],"exampleFix":"// before\nnew_def = PyTreeDef.deserialize_using_proto(registry, arbitrary_bytes)\n\n// after\nassert isinstance(data, bytes) and data, 'expect serialized pytree proto'\nnew_def = PyTreeDef.deserialize_using_proto(registry, data)  # from same jaxlib version","handlingStrategy":"validation","validationCode":"import jaxlib\nfrom jaxlib import pytree\n# only feed payloads from the same version\ndef safe_deserialize(registry, data: bytes):\n    if not isinstance(data, (bytes, bytearray)) or not data:\n        raise ValueError('empty/invalid pytree payload')\n    return pytree.PyTreeDef.deserialize_using_proto(registry, bytes(data))","typeGuard":null,"tryCatchPattern":"try:\n    treedef = pytree.PyTreeDef.deserialize_using_proto(registry, data)\nexcept Exception as e:\n    if 'Malformed pytree proto' in str(e):\n        raise ValueError(f'stal/corrupt pytree payload (jaxlib {jaxlib.__version__})') from e\n    raise","preventionTips":["Pin identical jax/jaxlib versions on producer and consumer","Treat serialized pytrees as version-locked artifacts; re-export after upgrades","Checksum payloads before transfer"],"tags":["jax","pytree","serialization","protobuf"],"backgroundTag":"deserialization-schema-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}