update docs and tests workflows to install jax dependencies

Also add jax identity example to the docs build

Author: pattonw

.github/workflows/docs.yaml

@@ -31,10 +31,11 @@ jobs:
 
       - name: Turn script into notebook
         run: |
-          uv sync --extra docs
+          uv sync --extra docs --extra torch --extra jax
           uv pip install .
           uv run python -m ipykernel install --name volara_env --user
           uv run jupytext --to notebook ./examples/cremi/cremi.py
+          uv run jupytext --to notebook ./examples/jax/identity.py
           cp -r examples docs/source
 
       - name: Build documentation

.github/workflows/tests.yaml

@@ -37,7 +37,7 @@ jobs:
 
     - name: Test with pytest
       run: |
-        uv sync --extra dev --resolution ${{ matrix.resolution }} --python ${{ matrix.python-version}}
+        uv sync --extra dev --extra torch --extra jax --resolution ${{ matrix.resolution }} --python ${{ matrix.python-version}}
         uv run pytest --cov=src tests
 
     - name: Upload coverage to Codecov

docs/source/index.rst

@@ -5,6 +5,7 @@
    overview
    install
    examples/cremi/cremi
+   examples/jax/identity
    api
 
 .. include:: ../../README.md