Merge pull request #598 from mrava87/doc-jax

doc: added jax to installation.rst

docs/source/installation.rst

@@ -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>`__.

requirements-doc.txt

@@ -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
@@ -21,6 +22,7 @@ docutils<0.18
 Sphinx
 pydata-sphinx-theme
 sphinx-gallery
+sphinxemoji
 numpydoc
 nbsphinx
 image