Skip to content

Add basic JAX support #84

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 20 commits into from
Feb 9, 2024
Merged

Add basic JAX support #84

Changes from 1 commit
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
2ee6902
Make is_numpy_array, is_cupy_array, is_torch_array, and is_dask_array…
asmeurer Feb 6, 2024
12b5294
Note dask in a couple of places in the README
asmeurer Feb 6, 2024
6bf5dad
Add JAX support
asmeurer Feb 6, 2024
583f6bb
Install JAX on CI
asmeurer Feb 6, 2024
9c8bed6
Install JAX as jax[cpu]
asmeurer Feb 6, 2024
6d59ae8
Fix a comment
asmeurer Feb 6, 2024
ce07cd9
Use jax.Array instead of jax.numpy.ndarray in is_jax_array
asmeurer Feb 6, 2024
ddb313e
Use the native jax to_device() method
asmeurer Feb 6, 2024
6004b97
Add a basic test for device() and to_device()
asmeurer Feb 6, 2024
701a5ef
XFAIL the device test for dask.array
asmeurer Feb 6, 2024
049d557
Allow to_device(x, "cpu") for JAX arrays
asmeurer Feb 6, 2024
db667ea
Move tests in test_common.py to test_helpers.py
asmeurer Feb 6, 2024
aafbbaa
Remove "cpu" device from jax to_device()
asmeurer Feb 8, 2024
919ec41
Merge branch 'main' into jax
asmeurer Feb 8, 2024
fa758f7
Rename import_or_skip_cupy to import_ and move jax logic into it
asmeurer Feb 8, 2024
6c338ca
Skip JAX tests in Python 3.8
asmeurer Feb 8, 2024
244462f
Rename test_helpers.py to test_common.py
asmeurer Feb 8, 2024
bff9bf2
Fix ruff warnings
asmeurer Feb 8, 2024
264e6c3
Add jax.numpy to the test_array_namespace()
asmeurer Feb 8, 2024
e7aff0f
Don't run the jax.numpy array_namespace test in Python 3.8
asmeurer Feb 8, 2024
File filter

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Rename test_helpers.py to test_common.py 
test_helpers.py is too confusing alongside tests/_helpers.py.
  • Loading branch information
@asmeurer
asmeurer committed Feb 8, 2024
commit 244462fd968a78b8ccf9e93bbde1bbe27566f0ce
0 tests/test_helpers.py → tests/test_common.py
View file
Open in desktop
File renamed without changes.