jax-ml/jax · error · XlaRuntimeError

Malformed pickled DictKey, expected 1-tuple

Error message

Malformed pickled DictKey, expected 1-tuple

What it means

DictKey.__setstate__ requires pickled state to be a 1-tuple holding the key object. A state tuple with more or fewer elements raises this error.

Source

Thrown at jaxlib/pytree.cc:1958

                     new (&key) SequenceKey(nb::cast<int>(state[0]));
                   });

  nb::class_<DictKey> dict_key(pytree, "DictKey",
                               nb::type_slots(DictKey::slots_),
                               nb::sig("class DictKey(typing.Hashable)"));
  dict_key.def(nb::init<nb::object>(), nb::arg("key").none());
  dict_key.def("__str__", &DictKey::ToString);
  dict_key.def("__repr__", &DictKey::ToReprString);
  dict_key.def("__eq__", &DictKey::Equals);
  dict_key.def("__hash__",
               [](const DictKey& key) { return nanobind::hash(key.key()); });
  dict_key.def_prop_ro("key", &DictKey::key);
  dict_key.def_prop_ro_static("__match_args__", &DictKey::MatchArgs);
  dict_key.def("__getstate__",
               [](DictKey& key) { return nb::make_tuple(key.key()); });
  dict_key.def("__setstate__", [](DictKey& key, const nb::tuple& state) {
    if (state.size() != 1) {
      throw xla::XlaRuntimeError("Malformed pickled DictKey, expected 1-tuple");
    }
    new (&key) DictKey(nb::cast<nb::object>(state[0]));
  });

  nb::class_<GetAttrKey> get_attr_key(
      pytree, "GetAttrKey", nb::sig("class GetAttrKey(typing.Hashable)"));
  get_attr_key.def(nb::init<nb::str>(), nb::arg("name"));
  get_attr_key.def("__str__", &GetAttrKey::ToString);
  get_attr_key.def("__repr__", &GetAttrKey::ToReprString);
  get_attr_key.def("__eq__", &GetAttrKey::Equals);
  get_attr_key.def("__hash__",
                   [](const GetAttrKey& key) { return nb::hash(key.name()); });
  get_attr_key.def_prop_ro("name", &GetAttrKey::name);
  get_attr_key.def_prop_ro_static("__match_args__", &GetAttrKey::MatchArgs);
  get_attr_key.def("__getstate__",
                   [](GetAttrKey& key) { return nb::make_tuple(key.name()); });
  get_attr_key.def("__setstate__", [](GetAttrKey& key, const nb::tuple& state) {
    if (state.size() != 1) {

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regenerate pickles with the current jaxlib
2. Ensure identical jaxlib versions when transferring pickled treedefs
3. Avoid mutating __getstate__ output

def valid_dict_key_state(s) -> bool:
    return isinstance(s, tuple) and len(s) == 1

When it happens

Trigger: Unpickling a jaxlib.pytree.DictKey whose state was altered or serialized by an incompatible version.

Common situations: Version-skewed pickles of treedef paths; manually assembled pickle payloads in tests or debugging tools.

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 PyTreeDef, expected 2-tuple

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