diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index c32638c..2e898e2 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -29,7 +29,7 @@ jobs:
       run: |
         sudo apt-get install -y libopenmpi-dev
         python -m pip install --upgrade pip
-        pip install jax==0.4.35
+        pip install jax
         pip install numpy setuptools cython wheel
         pip install git+https://github.com/MP-Gadget/pfft-python
         pip install git+https://github.com/MP-Gadget/pmesh