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.

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)

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 {aux_data=}
• CSR.tree_unflatten: invalid {aux_data=}
• CSC.tree_unflatten: invalid {aux_data=}
• Pytree serialization too large to deserialize.
• Could not deserialize PyTreeDefProto.

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