Moved jaxlib test dependency into CI scripts, so that the dev group doesn't assume the CPU.

Author: patrick-kidger

.github/workflows/release.yml

@@ -17,7 +17,7 @@ jobs:
                 cp -r ${{ github.workspace }}/tests ./tests
                 cp -r ${{ github.workspace }}/examples ./examples
                 cp ${{ github.workspace }}/pyproject.toml ./pyproject.toml
-                uv sync --extra dev --no-install-project --inexact
+                uv sync --extra dev --with 'jax[cpu]' --no-install-project --inexact
                 uv run --no-sync pytest
             pypi-token: ${{ secrets.pypi_token }}
             github-user: patrick-kidger

.github/workflows/run_tests.yml

@@ -44,7 +44,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install '.[dev]'
+          python -m pip install '.[dev]' 'jax[cpu]'
 
       - name: Test with pytest
         run: |

pyproject.toml

@@ -34,7 +34,6 @@ dependencies = [
 dev = [
     "pre-commit>=4.1.0",
     "beartype>=0.20.2",
-    "jax[cpu]",
     "pytest>=8.3.5",
     "pytest-env>=1.1.5",
 ]