diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index ff44938..682b1b6 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,4 +17,4 @@ jobs:
         python-version: 3.11
     - name: Test
       run: |
-        pip install absl-py graphviz && pip install . && python3 -m unittest discover -p "*_test.py"
+        pip install absl-py graphviz jax jaxlib && pip install . && python3 -m unittest discover -p "*_test.py"