force more recent jax

Author: lithomas1

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Install Dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest numpy torch dask[array] jax[cpu]
+          python -m pip install pytest numpy torch dask[array] jax[cpu]>=0.4.24
 
       - name: Run Tests
         run: |