fix(#177): jax version upgrade

fix(#177): jax version upgrade

Author: lgrcia

.github/workflows/ci.yml

@@ -15,11 +15,27 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - name: Set up Python 3.8
-        uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
         with:
-          python-version: 3.8
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install -U pip poetry
+          python -m poetry install --with dev
+      - name: Test with pytest
+        run: |
+          python -m poetry run python -m pytest
+
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
       - name: Install dependencies
         run: |
           python -m pip install -U pip poetry