4 Commits
Author SHA1 Message Date
Roman Knyazhitskiy 172b83b4fc Adding a test for generator support (#171)
* Add a test for generators

* Remove output annotations from decorators

Also guarded torch imports for better compatibility with
requirements.txt

* Add flag to the main meta class to skip the typecheck

* Return to the old solution

* Make async tests work

* Minor adjustments/fixing typos

* Correct Python path for new tests

* Remove some jax-dependent code

* Implement equality for MetaArrays

* Make all Dim variations frozen dataclasses

* Shorten AbstractArray methods

* Final touches

* Removing get_origin use

* Update tests with @jaxtyp
2024-02-25 12:07:01 +00:00
Patrick Kidger e03c1c329e 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.
2023-03-04 17:29:04 +00:00
Patrick Kidger 5600a1aac8 Fixed cloudpickle breaking, mark 2 2023-03-02 17:37:53 -08:00
Patrick Kidger 2b339715f9 Fixed cloudpickle breaking 2023-03-02 12:35:38 -08:00