Skip to content

Commit b199b3a

Browse files
authored
Merge pull request #598 from mrava87/doc-jax
doc: added jax to installation.rst
2 parents 21e590b + 628f60b commit b199b3a

File tree

2 files changed

+13
-1
lines changed

2 files changed

+13
-1
lines changed

docs/source/installation.rst

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -540,4 +540,14 @@ CuPy
540540
for GPU-accelerated computations. Since many different versions of CuPy exist (based on the
541541
CUDA drivers of the GPU), users must install CuPy prior to installing
542542
PyLops. To do so, follow their
543-
`installation instructions <https://docs.cupy.dev/en/stable/install.html>`__.
543+
`installation instructions <https://docs.cupy.dev/en/stable/install.html>`__.
544+
545+
546+
JAX
547+
---
548+
`JAX <http://JAX.readthedocs.io>`_ is another library that can be used as a drop-in replacement
549+
to NumPy and some parts of SciPy. It provides seamless support for multiple accelerators (e.g., GPUs, TPUs),
550+
Just-In-Time (JIT) compilation via Open XLA, and Automatic Differentiation. Similar to CuPy, since many
551+
different versions of JAX exist (based on the CUDA drivers of the GPU), users must install JAX prior
552+
to installing PyLops. To do so, follow their
553+
`installation instructions <https://jax.readthedocs.io/en/latest/installation.html#install-cpu>`__.

requirements-doc.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# since the dtcwt and spgl1 are not yet compatible with numpy=2.0.0
33
numpy>=1.21.0,<2.0.0
44
scipy>=1.11.0
5+
jax
56
--extra-index-url https://download.pytorch.org/whl/cpu
67
torch>=1.2.0
78
numba
@@ -21,6 +22,7 @@ docutils<0.18
2122
Sphinx
2223
pydata-sphinx-theme
2324
sphinx-gallery
25+
sphinxemoji
2426
numpydoc
2527
nbsphinx
2628
image

0 commit comments

Comments
 (0)