Backport PR #3184 on branch 1.2.x (ci: upper bound jax version due to…

… mrvi slowdown) (#3185)

Backport PR #3184: ci: upper bound jax version due to mrvi slowdown

Co-authored-by: Ori Kronfeld <ori.kronfeld@weizmann.ac.il>

pyproject.toml

@@ -35,8 +35,8 @@ dependencies = [
     "anndata>=0.11",
     "docrep>=0.3.2",
     "flax",
-    "jax",
-    "jaxlib",
+    "jax<=0.4.35",
+    "jaxlib<=0.4.35",
     "lightning>=2.0",
     "ml-collections>=0.1.1",
     "mudata>=0.1.2",
@@ -62,7 +62,7 @@ tests = ["pytest", "pytest-pretty", "coverage", "scvi-tools[optional]"]
 editing = ["jupyter", "pre-commit"]
 dev = ["scvi-tools[editing,tests]"]
 test = ["scvi-tools[tests]"]
-cuda = ["torchvision", "torchaudio", "jax[cuda]"]
+cuda = ["torchvision", "torchaudio", "jax[cuda]<=0.4.35"]
 metal = ["torchvision", "torchaudio", "jax-metal"]
 
 docs = [