jax-ml/jax · error · invalid_argument

Malformed pytree proto (invalid node type)

Error message

Malformed pytree proto (invalid node type)

What it means

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.

Solutions

1. Verify the payload was produced by the matching jaxlib version's serialize_using_proto on both ends
2. Upgrade/downgrade both producer and consumer to the same jaxlib/jax version
3. Round-trip test: serialize a simple pytree and deserialize it to validate the pipeline
4. If persisting pytrees, re-export them after upgrading JAX

Example fix

// before
new_def = PyTreeDef.deserialize_using_proto(registry, arbitrary_bytes)

// after
assert isinstance(data, bytes) and data, 'expect serialized pytree proto'
new_def = PyTreeDef.deserialize_using_proto(registry, data)  # from same jaxlib version

import jaxlib
from jaxlib import pytree
# only feed payloads from the same version
def safe_deserialize(registry, data: bytes):
    if not isinstance(data, (bytes, bytearray)) or not data:
        raise ValueError('empty/invalid pytree payload')
    return pytree.PyTreeDef.deserialize_using_proto(registry, bytes(data))

try:
    treedef = pytree.PyTreeDef.deserialize_using_proto(registry, data)
except Exception as e:
    if 'Malformed pytree proto' in str(e):
        raise ValueError(f'stal/corrupt pytree payload (jaxlib {jaxlib.__version__})') from e
    raise

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

• Parsing and encoding errors: unexpected token, malformed input — why parsers reject input and how to find the real culprit.

Related errors

• COO.tree_unflatten: invalid
• Could not deserialize PyTreeDefProto.
• CSC.tree_unflatten: invalid
• CSR.tree_unflatten: invalid
• Pytree serialization too large to deserialize.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/75028fbfa7a329a7. Report an issue: GitHub.

Appendix: source

Thrown at jaxlib/pytree.cc:1561

        break;
      case PyTreeNodeType::PY_TREE_KIND_NONE:
        node.kind = PyTreeKind::kNone;
        break;
      case PyTreeNodeType::PY_TREE_KIND_TUPLE:
        node.kind = PyTreeKind::kTuple;
        break;
      case PyTreeNodeType::PY_TREE_KIND_DICT:
        node.kind = PyTreeKind::kDict;
        for (uint32_t str_id : node_proto.dict_keys().str_id()) {
          if (str_id >= interned_strings.size()) {
            throw std::invalid_argument(
                "Malformed pytree proto (dict_key out of range).");
          }
          node.sorted_dict_keys.push_back(interned_strings.at(str_id));
        }
        break;
      default:
        throw std::invalid_argument(
            "Malformed pytree proto (invalid node type)");
        break;
    }
  }
  result->SetNumLeavesAndNumNodes();
  return result;
}

std::optional<std::pair<nb::object, nb::object>> PyTreeDef::GetNodeData()
    const {
  if (traversal_.empty()) {
    throw std::logic_error("empty PyTreeDef traversal.");
  }
  auto builtin_type = [](PyTypeObject* type_obj) {
    return nb::borrow<nb::object>(reinterpret_cast<PyObject*>(type_obj));
  };
  const auto& node = traversal_.back();
  switch (node.kind) {

View on GitHub (pinned to 1e1c6a8fc0)