jax-ml/jax · error · XlaRuntimeError

Could not deserialize PyTreeDefProto.

Error message

Could not deserialize PyTreeDefProto.

What it means

deserialize_using_proto failed at protobuf parsing: the bytes are not a valid PyTreeDefProto encoding. This is a parse-level failure distinct from semantic validation errors like 'invalid node type'.

Source

Thrown at jaxlib/pytree.cc:1887

    return AbslHashToPythonHash(absl::HashOf(t));
  });
  treedef.def("serialize_using_proto", [](const PyTreeDef& a) {
    PyTreeDefProto result;
    a.SerializeTo(result);
    std::string serialized = result.SerializeAsString();
    return nb::bytes(serialized.data(), serialized.size());
  });
  treedef.def_static(
      "deserialize_using_proto",
      [](nb_class_ptr<PyTreeRegistry> registry, nb::bytes data) {
        PyTreeDefProto input;
        std::string_view serialized(data.c_str(), data.size());
        if (serialized.size() > std::numeric_limits<int>::max()) {
          throw xla::XlaRuntimeError(
              "Pytree serialization too large to deserialize.");
        }
        if (!input.ParseFromArray(serialized.data(), serialized.size())) {
          throw xla::XlaRuntimeError("Could not deserialize PyTreeDefProto.");
        }
        return PyTreeDef::DeserializeFrom(std::move(registry), input);
      },
      nb::arg("registry"), nb::arg("data"));
  treedef.def("node_data", &PyTreeDef::GetNodeData,
              "Returns None if a leaf-pytree, else (type, node_data)",
              nb::sig("def node_data(self) -> tuple[type, Any] | None"));
  treedef.def_static(
      "from_node_data_and_children", &PyTreeDef::FromNodeDataAndChildren,
      nb::arg("registry"), nb::arg("node_data").none(), nb::arg("children"),
      "Reconstructs a pytree from `node_data()` and `children()`.",
      nb::sig(
          // clang-format off
        "def from_node_data_and_children("
        "self, "
        "registry: PyTreeRegistry, "
        "node_data: tuple[type, Any] | None, "
        "children: typing.Iterable[PyTreeDef]"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Confirm data came from serialize_using_proto of the same jaxlib version
2. Validate the bytes are non-empty and were transferred intact (checksum/length)
3. Match jax/jaxlib versions across producer and consumer

def is_likely_proto(data: bytes) -> bool:
    return isinstance(data, (bytes, bytearray)) and len(data) > 0

try:
    treedef = pytree.PyTreeDef.deserialize_using_proto(registry, data)
except Exception as e:
    if 'Could not deserialize PyTreeDefProto' in str(e):
        # regenerate payload with current jaxlib and retry once
        data = reserialize_from_source()
        treedef = pytree.PyTreeDef.deserialize_using_proto(registry, data)
    else:
        raise

When it happens

Trigger: Calling PyTreeDef.deserialize_using_proto(registry, data) where data is truncated, corrupted, not proto-encoded, or encoded with an incompatible proto schema.

Common situations: Passing JSON/pickle/str instead of proto bytes; truncated network transfer; version mismatch between the jaxlib that serialized and the one deserializing; passing the wrong field of a (treedef, leaves) tuple.

Related errors

• Malformed pytree proto (invalid node type)
• Pytree serialization too large to deserialize.
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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