jax-ml/jax · error · XlaRuntimeError

Malformed pickled PyTreeDef, expected 2-tuple

Error message

Malformed pickled PyTreeDef, expected 2-tuple

What it means

PyTreeDef.__setstate__ expects the pickled state to be a 2-tuple of (registry, treedef_data). If the pickle contains anything other than exactly 2 elements, this XlaRuntimeError is thrown.

Source

Thrown at jaxlib/pytree.cc:1913

  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]"
        ") -> PyTreeDef"
          // clang-format on
          ));
  treedef.def("__getstate__", &PyTreeDef::ToPickle);
  treedef.def("__setstate__", [](PyTreeDef& t, nb::object o) {
    nb::tuple pickle = nb::cast<nb::tuple>(o);
    if (pickle.size() != 2) {
      throw xla::XlaRuntimeError(
          "Malformed pickled PyTreeDef, expected 2-tuple");
    }
    auto registry = nb::cast<nb_class_ptr<PyTreeRegistry>>(pickle[0]);
    new (&t) PyTreeDef(registry);
    t.FromPickle(pickle[1]);
  });

  nb::class_<SequenceKey> sequence_key(
      pytree, "SequenceKey", nb::sig("class SequenceKey(typing.Hashable)"));
  sequence_key.def(nb::init<int>(), nb::arg("idx"));
  sequence_key.def("__str__", &SequenceKey::ToString);
  sequence_key.def("__repr__", &SequenceKey::ToReprString);
  sequence_key.def("__eq__", &SequenceKey::Equals);
  sequence_key.def("__hash__", [](const SequenceKey& key) {
    return key.idx() + kSequenceKeyHashSalt;
  });
  sequence_key.def_prop_ro("idx", &SequenceKey::idx);
  sequence_key.def_prop_ro_static("__match_args__", &SequenceKey::MatchArgs);

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-create the pickle with the current jaxlib version instead of loading stale ones
2. Verify the pickle contains a genuine PyTreeDef, not a modified tuple
3. Match JAX versions between save and load environments

import pickle
def valid_pytreedef_state(o) -> bool:
    return isinstance(o, tuple) and len(o) == 2

When it happens

Trigger: Unpickling a PyTreeDef whose __getstate__ output was tampered with, produced by a different jaxlib version with a different pickle layout, or manually constructed.

Common situations: Loading old pickles saved with an incompatible JAX version; hand-editing or wrapping pickled state; corrupted pickle files.

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

• The names should be exclusive and should not intersect in `n
• stop_gradient only works on valid JAX arrays, but input argu
• {self.__class__.__name__} has no attribute {name}
• Could not find type: %s.
• Malformed pickled SequenceKey, expected 1-tuple

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