diff --git a/pyproject.toml b/pyproject.toml
index 16d1091..2a4fb0e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,10 +30,11 @@ dependencies = [
     "pandas<2.0.0,>=1.4.0",
     "bbknn>1.5.0",
     "celltypist>1.3.0",
-    "chex<0.1.3",
     "gdown>4.6.0",
     "h5py>3.7.0",
     "huggingface-hub>=0.11.1",
+    "jax==0.4.19", # Remove with scvi-tools 1.1
+    "jaxlib==0.4.19",
     "onclass>=1.3",
     "scanorama>1.7.2",
     "scanpy>1.9.0",