jax-ml/jax · error · invalid_argument

Invalid StrongLRUCache pickle version, got %d, expected 1

Error message

Invalid StrongLRUCache pickle version, got %d, expected 1

What it means

StrongLRUCache.__setstate__ only accepts pickle version 1. Unpickling a cache whose pickled 'version' field differs (older or newer format) throws std::invalid_argument.

Source

Thrown at jaxlib/strong_lru_cache.cc:714

          .def_prop_ro("__wrapped__", &StrongLRUCache::wrapped)
          .def(
              "__getstate__",
              [](const StrongLRUCache& cache) {
                nb::dict pickle;
                pickle["version"] = 1;
                pickle["fn"] = cache.wrapped();
                pickle["cache_context_fn"] = cache.cache_context_fn();
                pickle["maxsize"] = cache.maxsize();
                pickle["explain"] = cache.explain();
                pickle["num_shards"] = cache.num_shards();
                return pickle;
              },
              nb::lock_self())
          .def("__setstate__", [](StrongLRUCache* cache,
                                  const nb::dict& pickle) {
            int version = nb::cast<int>(pickle["version"]);
            if (version != 1) {
              throw std::invalid_argument(absl::StrFormat(
                  "Invalid StrongLRUCache pickle version, got %d, expected 1",
                  version));
            }
            auto fn = nb::cast<nb::callable>(pickle["fn"]);
            auto cache_context_fn = nb::cast<std::optional<nb::callable>>(
                pickle["cache_context_fn"]);
            int64_t maxsize = nb::cast<int64_t>(pickle["maxsize"]);
            auto explain =
                nb::cast<std::optional<nb::callable>>(pickle["explain"]);
            int64_t num_shards = nb::cast<int64_t>(pickle["num_shards"]);

            new (cache)
                StrongLRUCache(std::move(cache_context_fn), std::move(fn),
                               maxsize, std::move(explain), num_shards);
          });

  strong_lru_cache.attr("__call__") = nb::steal<nb::object>(
      PyDescr_NewMethod(reinterpret_cast<PyTypeObject*>(strong_lru_cache.ptr()),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clear/regenerate the pickled cache with the current JAX version
2. Pin JAX versions when persisting cached objects
3. Avoid pickling internal LRU caches; persist the underlying data instead

state = pickle.loads(payload) if isinstance(payload, bytes) else payload
if not isinstance(state, dict) or state.get('version') != 1:
    raise ValueError('stale cache pickle; regenerate with current jaxlib')

try:
    cache = pickle.loads(payload)
except ValueError as e:
    if 'pickle version' in str(e):
        cache = rebuild_cache()  # cold start
    else:
        raise

When it happens

Trigger: Unpickling a jaxlib StrongLRUCache (used by jax._src.util.weakref_lru_cache / compiled-function caches) produced by a jaxlib version with a different pickle format version.

Common situations: Loading cached artifacts across JAX version upgrades/downgrades; stale cache files restored after an upgrade.

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}
• The error occurred in the __reduce__ method, which may indic
• key cannot be empty

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