We now have Float[np.ndarray, ...] <: np.ndarray. Added basic torch tests. (#68)

This required quite a lot of refactoring! JAX supports virtual subclass registration (its metaclass is ABCMeta) but NumPy does not, so we have to actually subclass `np.ndarray`.
Simple stuff like __base__ hacking fails due to deallocator conflicts.
This commit is contained in:
Patrick Kidger
2023-03-04 17:29:04 +00:00
committed by GitHub
parent fef81cf0a0
commit e03c1c329e
7 changed files with 252 additions and 283 deletions
+6
View File
@@ -1,4 +1,6 @@
import cloudpickle
import numpy as np
import torch
from jaxtyping import AbstractArray, Array, Shaped
@@ -6,5 +8,9 @@ from jaxtyping import AbstractArray, Array, Shaped
def test_pickle():
x = cloudpickle.dumps(Shaped[Array, ""])
y = cloudpickle.dumps(AbstractArray)
z = cloudpickle.dumps(Shaped[np.ndarray, ""])
w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
cloudpickle.loads(x)
cloudpickle.loads(y)
cloudpickle.loads(z)
cloudpickle.loads(w)