Update GitHub action unittest matrix

Adds a JAX nightly version and Python 3.13 and 3.14 to the version
lists.

Author: lukebaumann

.github/workflows/test.yml

@@ -23,8 +23,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.10', '3.11', '3.12']
-        jax-version: ['0.5.3', '0.6.0', '0.6.1', '0.6.2', '0.7.0', '0.7.1', '0.7.2']
+        python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
+        jax-version: ['0.5.3', '0.6.0', '0.6.1', '0.6.2', '0.7.0', '0.7.1', '0.7.2', 'nightly']
 
         exclude:
           - python-version: '3.10'
@@ -33,6 +33,8 @@ jobs:
             jax-version: '0.7.1'
           - python-version: '3.10'
             jax-version: '0.7.2'
+          - python-version: '3.10'
+            jax-version: 'nightly'
     runs-on: ubuntu-latest
     timeout-minutes: 30
 
@@ -48,7 +50,10 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - run: pip --version
-    - run: pip install -e . "jax==${{ matrix.jax-version }}"
+    - if: ${{ matrix.jax-version == 'nightly' }}
+      run: pip install -e . -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
+    - if: ${{ matrix.jax-version != 'nightly' }}
+      run: pip install -e . "jax==${{ matrix.jax-version }}"
     - run: pip freeze
 
     # Run tests