diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 9325def..83abbd7 100644 --- a/jaxtyping/array_types.py +++ b/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): diff --git a/test/requirements.txt b/test/requirements.txt index dd54a1e..2492cb3 100644 --- a/test/requirements.txt +++ b/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 diff --git a/test/test_serialisation.py b/test/test_serialisation.py new file mode 100644 index 0000000..b1d0e54 --- /dev/null +++ b/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)