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 with the integer flattened index. Any other tuple length is rejected.

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)

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

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