jax-ml/jax · error · XlaRuntimeError

Pytree serialization too large to deserialize.

Error message

Pytree serialization too large to deserialize.

What it means

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.

Source

Thrown at jaxlib/pytree.cc:1883

  treedef.def("__ne__", [](const PyTreeDef& a, nb::object b) {
    return nb::isinstance<PyTreeDef>(b) && a != nb::cast<PyTreeDef>(b);
  });
  treedef.def("__hash__", [](const PyTreeDef& t) -> Py_hash_t {
    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("

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the pytree size: flatten into leaves + treedef separately, or chunk the structure
2. Check for accidental double serialization (bytes of bytes) inflating size
3. Avoid serializing gigantic static structures as pytrees; use arrays

import sys
MAX = 2**31 - 1
def safe_deserialize(registry, data: bytes):
    if len(data) > MAX:
        raise ValueError(f'payload {len(data)} exceeds protobuf int limit')
    return pytree.PyTreeDef.deserialize_using_proto(registry, data)

When it happens

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

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

Related errors

• Malformed pytree proto (invalid node type)
• Could not deserialize PyTreeDefProto.
• 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/2e0c34aacd53b986. Report an issue: GitHub.