diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e5060bf..a93baef 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -50,9 +50,9 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - run: pip --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/
+      run: pip install -e ".[test]" -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==${JAX_VERSION}"
+      run: pip install -e ".[test]" "jax==${JAX_VERSION}"
       env:
         JAX_VERSION: ${{ matrix.jax-version }}
     - run: pip freeze