Fixed working with the new (unreleased) version of typeguard (#53)

README.md

@@ -28,8 +28,6 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires JAX 0.3.4+.
-
 Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
 
 ## Documentation

jaxtyping/__init__.py

@@ -89,4 +89,4 @@ elif has_jax:
 
 del has_jax
 
-__version__ = "0.2.9"
+__version__ = "0.2.10"

jaxtyping/pytree_type.py

@@ -59,6 +59,14 @@ class _MetaPyTree(type):
         return out
 
 
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
+
+
 class _MetaSubscriptPyTree(type):
     def __call__(self, *args, **kwargs):
         raise RuntimeError("PyTree cannot be instantiated")
@@ -77,7 +85,7 @@ class _MetaSubscriptPyTree(type):
         def is_leaftype(x):
             try:
                 accepts_leaftype(x)
-            except TypeError:
+            except _TypeCheckError:
                 return False
             else:
                 return True

test/helpers.py

@@ -18,16 +18,31 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import equinox as eqx
+import typeguard
 
 
+ParamError = []
+ReturnError = []
+ParamError.append(TypeError)  # old typeguard
+ReturnError.append(TypeError)  # old typeguard
+
+try:
+    # new typeguard
+    ParamError.append(typeguard.TypeCheckError)
+    ReturnError.append(typeguard.TypeCheckError)
+except AttributeError:
+    pass
+
 try:
     import beartype
 except ImportError:
-    ParamError = TypeError
+    pass
-    ReturnError = TypeError
 else:
-    ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation)
+    ParamError.append(beartype.roar.BeartypeCallHintParamViolation)
-    ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation)
+    ReturnError.append(beartype.roar.BeartypeCallHintReturnViolation)
+
+ParamError = tuple(ParamError)
+ReturnError = tuple(ReturnError)
 
 
 @eqx.filter_jit