From 5600a1aac8298f57ab34dc0cc5e988f66c432071 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Thu, 2 Mar 2023 17:37:53 -0800 Subject: [PATCH] Fixed cloudpickle breaking, mark 2 --- jaxtyping/array_types.py | 70 -------------------- test/test_array.py | 128 ++++++++++++++++++++++++++----------- test/test_serialisation.py | 6 +- 3 files changed, 95 insertions(+), 109 deletions(-) diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 83abbd7..595fcfd 100644 --- a/jaxtyping/array_types.py +++ b/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 diff --git a/test/test_array.py b/test/test_array.py index 585fd66..7cf8a8f 100644 --- a/test/test_array.py +++ b/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"], + }, + ) diff --git a/test/test_serialisation.py b/test/test_serialisation.py index b1d0e54..bb118af 100644 --- a/test/test_serialisation.py +++ b/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)