Include jax extra in badge workflow install

Switch the install step from .[test] to .[test,jax] so JAX-gated tests (those that use pytest.importorskip('jax') at module level) are collected and counted in the published badges.

Author: timlichtenberg

.github/workflows/publish-test-badges.yml

@@ -35,7 +35,7 @@ jobs:
       - name: Install package with test extra
         run: |
           python -m pip install --upgrade pip
-          pip install ".[test]"
+          pip install ".[test,jax]"
 
       - name: Generate badges
         run: python tools/generate_test_badges.py --out .github/badges/