Added support for custom arrays.

jaxtyping/_array_types.py

@@ -188,14 +188,12 @@ def __instancecheck_str__(cls, obj: Any) -> str:
             # TensorFlow
             dtype = obj.dtype.as_numpy_dtype.__name__
         else:
-            # PyTorch
-            repr_dtype = repr(obj.dtype).split(".")
-            if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
-                dtype = repr_dtype[1]
-            else:
-                raise AnnotationError(
-                    "Unrecognised array/tensor type to extract dtype from"
-                )
+            # Everyone else, including PyTorch.
+            # This offers an escape hatch for anyone looking to use jaxtyping for their
+            # own array-like types.
+            dtype = obj.dtype
+            if not isinstance(dtype, str):
+                *_, dtype = repr(obj.dtype).rsplit(".", 1)
 
         if cls.dtypes is not _any_dtype:
             in_dtypes = False

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.30"
+version = "0.2.31"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"

test/test_array.py

@@ -710,3 +710,59 @@ def test_scalar_variadic_dim():
 def test_scalar_dtype_mismatch():
     with pytest.raises(ValueError):
         Float[bool, "..."]
+
+
+def test_custom_array(jaxtyp, typecheck):
+    class MyArray1:
+        @property
+        def dtype(self):
+            return "foo"
+
+        @property
+        def shape(self):
+            return (3,)
+
+    class MyArray2:
+        @property
+        def dtype(self):
+            return "bar"
+
+        @property
+        def shape(self):
+            return (3,)
+
+    class MyArray3:
+        @property
+        def dtype(self):
+            return "foo"
+
+        @property
+        def shape(self):
+            return (4,)
+
+    class FooDtype(AbstractDtype):
+        dtypes = ["foo"]
+
+    @jaxtyp(typecheck)
+    def f(x: FooDtype[MyArray1, "3"]):
+        pass
+
+    f(MyArray1())
+    with pytest.raises(ParamError):
+        f(MyArray2())
+    with pytest.raises(ParamError):
+        f(MyArray3())
+
+    @jaxtyp(typecheck)
+    def g(x: FooDtype[MyArray1, "3"], y: FooDtype[MyArray1, "4"]):
+        pass
+
+    with pytest.raises(ParamError):
+        g(MyArray1(), MyArray1())
+
+    @jaxtyp(typecheck)
+    def h(x: FooDtype[MyArray1, "3"], y: FooDtype[MyArray3, "4"]):
+        pass
+
+    with pytest.raises(ParamError):
+        g(MyArray1(), MyArray3())