mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-21 11:16:29 +08:00
Patrick Kidger
e03c1c329e
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.
17 lines
412 B
Python
17 lines
412 B
Python
import cloudpickle
|
|
import numpy as np
|
|
import torch
|
|
|
|
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)
|