wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/run_tests.yml
+2 -1
View File
@@ -42,7 +42,8 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install pytest wheel beartype equinox jaxlib
python -m pip install pytest wheel beartype equinox jaxlib cloudpickle
python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
- name: Checks with pre-commit
uses: pre-commit/action@v2.0.3