Merge pull request #505 from FlorianPfaff/jaxnew3

Added jax backend

Author: FlorianPfaff

.github/workflows/mega-linter.yml

@@ -102,12 +102,13 @@ jobs:
         run: |
           sed 's/==.*//' requirements-dev.txt > requirements-dev_no_version.txt
 
-      - name: Remove torch, tirton, and nvidia entries (unsupported by alpine)
+      - name: Remove torch, triton, jax, and nvidia entries (unsupported by alpine)
         if: steps.cache-wheels.outputs.cache-hit != 'true'
         run: |
           sed -i '/^torch/d' requirements-dev_no_version.txt
           sed -i '/^nvidia/d' requirements-dev_no_version.txt
           sed -i '/^triton/d' requirements-dev_no_version.txt
+          sed -i '/^jax/d' requirements-dev_no_version.txt
 
       - name: Run CMake to find LAPACK
         if: steps.cache-wheels.outputs.cache-hit != 'true'

.github/workflows/tests.yml

@@ -79,14 +79,43 @@ jobs:
         env:
           PYTHONPATH: ${{ github.workspace }}
 
-      - name: Upload pytorch test result artifact
+  test-jax:
+    runs-on: ubuntu-latest
+    permissions:
+      checks: write
+      pull-requests: write
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install poetry
+          poetry env use python
+          poetry install --extras "healpy_support" --extras "jax_support"
+
+      - name: Run tests with jax backend
+        run: |
+          export PYRECEST_BACKEND=jax
+          poetry run python -m pytest --rootdir . -v --strict-config --junitxml=junit_test_results_jax.xml ./pyrecest
+        env:
+          PYTHONPATH: ${{ github.workspace }}
+
+      - name: Upload jax test result artifact
         uses: actions/upload-artifact@v3
         with:
-          name: pytorch-test-results
-          path: junit_test_results_pytorch.xml
+          name: jax-test-results
+          path: junit_test_results_jax.xml
 
   publish-results:
-    needs: [test-numpy, test-pytorch]
+    needs: [test-numpy, test-pytorch, test-jax]
     runs-on: ubuntu-latest
     if: always()
     permissions:
@@ -104,3 +133,4 @@ jobs:
           files: |
             test-results/numpy-test-results/junit_test_results_numpy.xml
             test-results/pytorch-test-results/junit_test_results_pytorch.xml
+            test-results/pytorch-test-results/junit_test_results_jax.xml

poetry.lock

@@ -3,7 +3,7 @@ name = "pyrecest"
 description = "Framework for recursive Bayesian estimation in Python."
 readme = "README.md"
 authors = ["Florian Pfaff <pfaff@kit.edu>"]
-version = "0.6.1"
+version = "0.7.0"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
@@ -20,10 +20,14 @@ shapely = "*"
 [tool.poetry.extras]
 healpy_support = ["healpy"]
 pytorch_support = ["torch"]
+jax_support = ["jax", "jaxlib", "autograd"]
 
 [tool.poetry.group.dev.dependencies]
 healpy = "*"
 torch = "*"
+jax = "*"
+jaxlib = "*"
+autograd = "*"
 autopep8 = "^2.0.2"
 pytest = "*"
 parameterized = "*"