jax-ml/jax · error · XlaRuntimeError

Malformed pickled SequenceKey, expected 1-tuple

Error message

Malformed pickled SequenceKey, expected 1-tuple

What it means

SequenceKey.__setstate__ requires its pickled state to be a 1-tuple containing the integer index. Any other tuple length is rejected as malformed.

Source

Thrown at jaxlib/pytree.cc:1937

  });

  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);
  sequence_key.def("__getstate__",
                   [](SequenceKey& key) { return nb::make_tuple(key.idx()); });
  sequence_key.def("__setstate__",
                   [](SequenceKey& key, const nb::tuple& state) {
                     if (state.size() != 1) {
                       throw xla::XlaRuntimeError(
                           "Malformed pickled SequenceKey, expected 1-tuple");
                     }
                     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()); });

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Don't hand-build pickle state; unpickle objects produced by the same jaxlib
2. Re-pickle via __getstate__ after upgrading
3. Upgrade both environments to matching jaxlib versions

def valid_seq_key_state(s) -> bool:
    return isinstance(s, tuple) and len(s) == 1 and isinstance(s[0], int)

When it happens

Trigger: Unpickling a jaxlib.pytree.SequenceKey (an entry in treedef node paths) whose pickled state tuple doesn't have exactly one element, e.g. crafted or version-skewed pickles.

Common situations: Manually constructing/reconstructing key objects; pickles produced by an older jaxlib that serialized different fields.

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/27f17d7e6c6af566. Report an issue: GitHub.