mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-09-11 12:21:38 +08:00
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user