Merge pull request #590 from mrava87/feature-jax

Feature: jax integration

Author: mrava87

docs/source/api/index.rst

@@ -29,6 +29,7 @@ Templates
     FunctionOperator
     MemoizeOperator
     TorchOperator
+    JaxOperator
 
 Basic operators
 ~~~~~~~~~~~~~~~

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),