Merge pull request #590 from mrava87/feature-jax

Feature: jax integration
PyLops · Aug 9, 2024 · 21e590b · 21e590b
2 parents 32bad36 + 054c5b8
commit 21e590b
Show file tree

Hide file tree

Showing 38 changed files with 1,434 additions and 362 deletions.
diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst
@@ -29,6 +29,7 @@ Templates
     FunctionOperator
     MemoizeOperator
     TorchOperator
+    JaxOperator
 
 Basic operators
 ~~~~~~~~~~~~~~~

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -21,6 +21,7 @@
     "numpydoc",
     "nbsphinx",
     "sphinx_gallery.gen_gallery",
+    "sphinxemoji.sphinxemoji",
     # 'sphinx.ext.napoleon',
 ]
 
@@ -29,6 +30,8 @@
     "python": ("https://docs.python.org/3/", None),
     "numpy": ("https://docs.scipy.org/doc/numpy/", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
+    "cupy": ("https://docs.cupy.dev/en/stable/", None),
+    "jax": ("https://jax.readthedocs.io/en/latest", None),
     "sklearn": ("http://scikit-learn.org/stable/", None),
     "pandas": ("http://pandas.pydata.org/pandas-docs/stable/", None),
     "matplotlib": ("https://matplotlib.org/", None),