Fixed cloudpickle breaking, mark 2

jaxtyping/array_types.py

@@ -67,72 +67,24 @@ class _NamedDim:
         self.name = name
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.name != other.name:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.name, self.broadcastable))
-
 
 class _NamedVariadicDim:
     def __init__(self, name, broadcastable):
         self.name = name
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.name != other.name:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.name, self.broadcastable))
-
 
 class _FixedDim:
     def __init__(self, size, broadcastable):
         self.size = size
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.size != other.size:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.size, self.broadcastable))
-
 
 class _SymbolicDim:
     def __init__(self, expr, broadcastable):
         self.expr = expr
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.expr != other.expr:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.expr, self.broadcastable))
-
 
 _AbstractDimOrVariadicDim = Union[
     Literal[_anonymous_dim],
@@ -184,28 +136,6 @@ def _check_dims(
 
 
 class _MetaAbstractArray(type):
-    def __eq__(self, other):
-        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):
-        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):
             return False

test/test_array.py

@@ -24,7 +24,16 @@ import jax.random as jr
 import numpy as np
 import pytest
 
-from jaxtyping import AbstractDtype, Array, ArrayLike, Float, Float32, jaxtyped, Shaped
+from jaxtyping import (
+    AbstractArray,
+    AbstractDtype,
+    Array,
+    ArrayLike,
+    Float,
+    Float32,
+    jaxtyped,
+    Shaped,
+)
 
 from .helpers import ParamError, ReturnError
 
@@ -414,6 +423,36 @@ def test_incomplete_symbolic(typecheck, getkey):
         foo(x)
 
 
+def _eq(x, y):
+    assert type(x) is set
+    assert type(y) is set
+    assert len(x) == len(y)
+    for xi in x:
+        assert any(_eq_impl(xi, yi) for yi in y)
+
+
+def _eq_impl(x, y):
+    if issubclass(x, AbstractArray):
+        if type(x) is not type(y):
+            return False
+        if x.array_type is not y.array_type:
+            return False
+        if x.dtypes != y.dtypes:
+            return False
+        if x.index_variadic != y.index_variadic:
+            return False
+        if len(x.dims) != len(y.dims):
+            return False
+        for x_dim, y_dim in zip(x.dims, y.dims):
+            if type(x_dim) is not type(y_dim):
+                return False
+            if x_dim.__dict__ != y_dim.__dict__:
+                return False
+        return True
+    else:
+        return x is y
+
+
 def test_arraylike(typecheck, getkey):
     floatlike1 = Float32[ArrayLike, ""]
     floatlike2 = Float[ArrayLike, ""]
@@ -422,44 +461,59 @@ def test_arraylike(typecheck, getkey):
     assert get_origin(floatlike1) is Union
     assert get_origin(floatlike2) is Union
     assert get_origin(floatlike3) is Union
-    assert set(get_args(floatlike1)) == {
-        Float32[Array, ""],
-        Float32[np.ndarray, ""],
-        Float32[np.bool_, ""],
-        Float32[np.number, ""],
-        float,
-    }
-    assert set(get_args(floatlike2)) == {
-        Float[Array, ""],
-        Float[np.ndarray, ""],
-        Float[np.bool_, ""],
-        Float[np.number, ""],
-        float,
-    }
-    assert set(get_args(floatlike3)) == {
-        Float32[Array, "4"],
-        Float32[np.ndarray, "4"],
-        Float32[np.bool_, "4"],
-        Float32[np.number, "4"],
-    }
+    _eq(
+        set(get_args(floatlike1)),
+        {
+            Float32[Array, ""],
+            Float32[np.ndarray, ""],
+            Float32[np.bool_, ""],
+            Float32[np.number, ""],
+            float,
+        },
+    )
+    _eq(
+        set(get_args(floatlike2)),
+        {
+            Float[Array, ""],
+            Float[np.ndarray, ""],
+            Float[np.bool_, ""],
+            Float[np.number, ""],
+            float,
+        },
+    )
+    _eq(
+        set(get_args(floatlike3)),
+        {
+            Float32[Array, "4"],
+            Float32[np.ndarray, "4"],
+            Float32[np.bool_, "4"],
+            Float32[np.number, "4"],
+        },
+    )
 
     shaped1 = Shaped[ArrayLike, ""]
     shaped2 = Shaped[ArrayLike, "4"]
     assert get_origin(shaped1) is Union
     assert get_origin(shaped2) is Union
-    assert set(get_args(shaped1)) == {
-        Shaped[Array, ""],
-        Shaped[np.ndarray, ""],
-        Shaped[np.bool_, ""],
-        Shaped[np.number, ""],
-        bool,
-        int,
-        float,
-        complex,
-    }
-    assert set(get_args(shaped2)) == {
-        Shaped[Array, "4"],
-        Shaped[np.ndarray, "4"],
-        Shaped[np.bool_, "4"],
-        Shaped[np.number, "4"],
-    }
+    _eq(
+        set(get_args(shaped1)),
+        {
+            Shaped[Array, ""],
+            Shaped[np.ndarray, ""],
+            Shaped[np.bool_, ""],
+            Shaped[np.number, ""],
+            bool,
+            int,
+            float,
+            complex,
+        },
+    )
+    _eq(
+        set(get_args(shaped2)),
+        {
+            Shaped[Array, "4"],
+            Shaped[np.ndarray, "4"],
+            Shaped[np.bool_, "4"],
+            Shaped[np.number, "4"],
+        },
+    )

test/test_serialisation.py

@@ -4,5 +4,7 @@ from jaxtyping import AbstractArray, Array, Shaped
 
 
 def test_pickle():
-    cloudpickle.dumps(Shaped[Array, ""])
-    cloudpickle.dumps(AbstractArray)
+    x = cloudpickle.dumps(Shaped[Array, ""])
+    y = cloudpickle.dumps(AbstractArray)
+    cloudpickle.loads(x)
+    cloudpickle.loads(y)