Skip to content
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

doc: added jax to installation.rst #598

Merged
merged 3 commits into from
Aug 9, 2024
Merged

doc: added jax to installation.rst #598

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ad6c6f3
doc: added jax to installation.rst
mrava87 Aug 9, 2024
fae0f00
bug: added sphinxemoji to requirements-doc
mrava87 Aug 9, 2024
628f60b
bug: added jax to requirements-doc
mrava87 Aug 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion docs/source/installation.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -540,4 +540,14 @@ CuPy
for GPU-accelerated computations. Since many different versions of CuPy exist (based on the
CUDA drivers of the GPU), users must install CuPy prior to installing
PyLops. To do so, follow their
`installation instructions <https://docs.cupy.dev/en/stable/install.html>`__.
`installation instructions <https://docs.cupy.dev/en/stable/install.html>`__.


JAX
---
`JAX <http://JAX.readthedocs.io>`_ is another library that can be used as a drop-in replacement
to NumPy and some parts of SciPy. It provides seamless support for multiple accelerators (e.g., GPUs, TPUs),
Just-In-Time (JIT) compilation via Open XLA, and Automatic Differentiation. Similar to CuPy, since many
different versions of JAX exist (based on the CUDA drivers of the GPU), users must install JAX prior
to installing PyLops. To do so, follow their
`installation instructions <https://jax.readthedocs.io/en/latest/installation.html#install-cpu>`__.
2 changes: 2 additions & 0 deletions requirements-doc.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# since the dtcwt and spgl1 are not yet compatible with numpy=2.0.0
numpy>=1.21.0,<2.0.0
scipy>=1.11.0
jax
--extra-index-url https://download.pytorch.org/whl/cpu
torch>=1.2.0
numba
Expand All @@ -21,6 +22,7 @@ docutils<0.18
Sphinx
pydata-sphinx-theme
sphinx-gallery
sphinxemoji
numpydoc
nbsphinx
image
Expand Down