diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 24b0814..0f6f98a 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -29,7 +29,6 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
-        python -m pip install jaxlib
         # install the package locally
         python -m pip install .
     - name: Lint with flake8
diff --git a/pyproject.toml b/pyproject.toml
index d2042cb..66aee49 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,8 @@ dependencies = [
     "numpy",
     "tqdm",
     "xarray",
-    "jax"
+    "jax",
+    "jaxlib"
 ]
 dynamic = ["version"]