Fixed cloudpickle breaking

jaxtyping/array_types.py

@@ -185,20 +185,26 @@ def _check_dims(
 
 class _MetaAbstractArray(type):
     def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.array_type is not other.array_type:
-            return False
-        if self.dtypes != other.dtypes:
-            return False
-        if self.dims != other.dims:
-            return False
-        if self.index_variadic != other.index_variadic:
-            return False
-        return True
+        if self is AbstractArray:
+            return other is AbstractArray
+        else:
+            if type(self) is not type(other):
+                return False
+            if self.array_type is not other.array_type:
+                return False
+            if self.dtypes != other.dtypes:
+                return False
+            if self.dims != other.dims:
+                return False
+            if self.index_variadic != other.index_variadic:
+                return False
+            return True
 
     def __hash__(self):
-        return hash((self.array_type, self.dtypes, self.dims, self.index_variadic))
+        if self is AbstractArray:
+            return 0
+        else:
+            return hash((self.array_type, self.dtypes, self.dims, self.index_variadic))
 
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):

test/requirements.txt

@@ -2,3 +2,4 @@ equinox>=0.5.3
 pytest>=7.0.1
 beartype>=0.10.4
 typeguard>=2.13.3
+cloudpickle>=2.2.1

test/test_serialisation.py

@@ -0,0 +1,8 @@
+import cloudpickle
+
+from jaxtyping import AbstractArray, Array, Shaped
+
+
+def test_pickle():
+    cloudpickle.dumps(Shaped[Array, ""])
+    cloudpickle.dumps(AbstractArray)