array shapes+dtypes no longer checked as part of pytree flattening. This avoids edge-case crash when using pytree-path dependent sizes

jaxtyping/_array_types.py

@@ -28,7 +28,12 @@ from typing import Any, Literal, NoReturn, Optional, Union
 import numpy as np
 
 from ._raise import jaxtyping_raise, jaxtyping_raise_from
-from ._storage import get_shape_memo, get_treepath_memo, set_shape_memo
+from ._storage import (
+    get_shape_memo,
+    get_treeflatten_memo,
+    get_treepath_memo,
+    set_shape_memo,
+)
 
 
 try:
@@ -160,6 +165,8 @@ class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
+        if get_treeflatten_memo():
+            return True
 
         if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy

jaxtyping/_pytree_type.py

@@ -26,9 +26,11 @@ import typeguard
 
 from ._raise import jaxtyping_raise_from
 from ._storage import (
+    clear_treeflatten_memo,
     clear_treepath_memo,
     get_shape_memo,
     set_shape_memo,
+    set_treeflatten_memo,
     set_treepath_memo,
 )
 
@@ -79,7 +81,7 @@ class _MetaPyTree(type):
             def is_flatten_leaftype(x):
                 return False
 
-            def is_check_leaftype(x, new_scope):
+            def is_check_leaftype(x):
                 return True
 
         else:
@@ -93,22 +95,21 @@ class _MetaPyTree(type):
             def accepts_leaftype(x: cls.leaftype):
                 pass
 
-            def is_leaftype(x, new_scope=True):
-                if new_scope and cls.structure is not None:
-                    set_treepath_memo(None, cls.structure)
+            def is_leaftype(x):
                 try:
                     accepts_leaftype(x)
                 except _TypeCheckError:
                     return False
                 else:
                     return True
-                finally:
-                    if new_scope and cls.structure is not None:
-                        clear_treepath_memo()
 
             is_flatten_leaftype = is_check_leaftype = is_leaftype
 
-        leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        set_treeflatten_memo()
+        try:
+            leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        finally:
+            clear_treeflatten_memo()
         if cls.structure is not None:
             if cls.structure.isidentifier():
                 try:
@@ -173,7 +174,7 @@ class _MetaPyTree(type):
             for leaf_index, leaf in enumerate(leaves):
                 if cls.structure is not None:
                     set_treepath_memo(leaf_index, cls.structure)
-                if not is_check_leaftype(leaf, new_scope=False):
+                if not is_check_leaftype(leaf):
                     return False
                 clear_treepath_memo()
         finally:

jaxtyping/_storage.py

@@ -104,3 +104,21 @@ def get_treepath_memo() -> str:
             )
         )
     return _treepath_storage.value
+
+
+_treeflatten_storage = threading.local()
+
+
+def clear_treeflatten_memo() -> None:
+    _treeflatten_storage.value = False
+
+
+def set_treeflatten_memo():
+    _treeflatten_storage.value = True
+
+
+def get_treeflatten_memo():
+    try:
+        return _treeflatten_storage.value
+    except AttributeError:
+        return False