jax-ml/jax · error · XlaRuntimeError

Malformed pickled FlattenedIndexKey, expected 1-tuple

Error message

Malformed pickled FlattenedIndexKey, expected 1-tuple

What it means

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

Solutions

1. Regenerate the pickled objects under the current jaxlib
2. Keep jaxlib versions aligned across processes

def valid_flattened_index_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.FlattenedIndexKey whose state tuple doesn't have exactly one element.

Common situations: Cross-version unpickling of treedef paths containing flattened index keys; manipulated pickle data.

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 DictKey, 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/abde9352a2c0943e. Report an issue: GitHub.

Appendix: source

Thrown at jaxlib/pytree.cc:2002

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

}  // namespace jax

View on GitHub (pinned to 1e1c6a8fc0)