Feature: jax integration

.github/workflows/build.yaml

@@ -7,7 +7,7 @@ jobs:
     strategy:
       matrix:
         platform: [ ubuntu-latest, macos-13 ]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11"]
 
     runs-on: ${{ matrix.platform }}
     steps:

azure-pipelines.yml

@@ -26,7 +26,7 @@ jobs:
 #   steps:
 #   - task: UsePythonVersion@0
 #     inputs:
-#       versionSpec: '3.7'
+#       versionSpec: '3.9'
 #       architecture: 'x64'
 #
 #   - script: |
@@ -55,7 +55,7 @@ jobs:
   steps:
   - task: UsePythonVersion@0
     inputs:
-      versionSpec: '3.8'
+      versionSpec: '3.9'
       architecture: 'x64'
 
   - script: |
@@ -84,7 +84,7 @@ jobs:
   steps:
   - task: UsePythonVersion@0
     inputs:
-      versionSpec: '3.8'
+      versionSpec: '3.9'
       architecture: 'x64'
 
   - script: |

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