WIP

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -401,7 +401,8 @@ def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompa
 
             id_ = id(obj)
             try:
-                return id_, seen[id_]
+                kind = seen[id_]
+                return None if kind is None else (id_, kind)
             except KeyError:
                 pass
 
@@ -487,19 +488,19 @@ class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL01,RT01
 
         @override
         def persistent_load(self, pid: tuple[int, int]) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
-            prev_id, kind = pid
             try:
-                return seen[prev_id]
+                return seen[pid]
             except KeyError:
                 pass
 
+            _, kind = pid
             try:
                 obj = next(iters[kind])
             except StopIteration as e:
                 msg = "Not enough objects to unpickle"
                 raise ValueError(msg) from e
 
-            seen[prev_id] = obj
+            seen[pid] = obj
             return obj
 
     f = io.BytesIO(pik)