jax-ml/jax · error · XlaRuntimeError

Malformed pickled DictKey, expected 1-tuple

Error message

Malformed pickled DictKey, expected 1-tuple

What it means

DictKey.__setstate__ expects a 1-tuple holding the key object. It fires when the pickled state is not that shape, i.e. the DictKey pickle is malformed.

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

• Could not find type: .
• Malformed pickled FlattenedIndexKey, expected 1-tuple
• Malformed pickled GetAttrKey, expected 1-tuple
• Malformed pickled PyTreeDef, expected 2-tuple
• Malformed pickled SequenceKey, expected 1-tuple

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

Appendix: 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)