ci: remove Jax version constraint (#3309)

Removing the jax constraint and adding a warning for version in mrVI
training, as jax <=0.4.35 is not supported by newer versions of CUDA
which becomes more popular and default

Author: ori-kron-wis

CHANGELOG.md

@@ -33,6 +33,8 @@ to [Semantic Versioning]. Full commit history is available in the
 
 #### Removed
 
+- Removed Jax version constraint for mrVI training. {pr}`3309`.
+
 ### 1.3.0 (2025-02-28)
 
 #### Added

docs/installation.md

@@ -49,11 +49,20 @@ scvi-tools depends on PyTorch and JAX for accelerated computing. If you don't pl
 an accelerated device, we recommend installing scvi-tools directly and letting these dependencies
 be installed automatically by your package manager of choice.
 
-If you plan on taking advantage of an accelerated device (e.g. Nvidia GPU or Apple Silicon), we
-recommend installing PyTorch and JAX _before_ installing scvi-tools. Please follow the respective
-installation instructions for [PyTorch](https://pytorch.org/get-started/locally/) and
-[JAX](https://jax.readthedocs.io/en/latest/installation.html) compatible with your system and
-device type.
+If you plan on taking advantage of an accelerated device (e.g. Nvidia GPU or Apple Silicon), scvi-tools supports it.
+In order to install scvi-tools with Nvidia GPU CUDA support use:
+```bash
+pip install -U scvi-tools[cuda]
+```
+And for Apple Silicon metal (MPS) support:
+```bash
+pip install -U scvi-tools[metal]
+```
+
+However, there might be cases where the GPU HW is not supporting the latest installation of PyTorch and Jax.
+In this case we recommend installing PyTorch and JAX _before_ installing scvi-tools.
+Please follow the respective installation instructions for [PyTorch](https://pytorch.org/get-started/locally/) and
+[JAX](https://jax.readthedocs.io/en/latest/installation.html) compatible with your system and device type.
 
 ## Optional dependencies
 

pyproject.toml

@@ -35,8 +35,8 @@ dependencies = [
     "anndata>=0.11",
     "docrep>=0.3.2",
     "flax",
-    "jax<=0.4.35",
-    "jaxlib<=0.4.35",
+    "jax",
+    "jaxlib",
     "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]<=0.4.35"]
+cuda = ["torchvision", "torchaudio", "jax[cuda12]"]
 metal = ["torchvision", "torchaudio", "jax-metal"]
 
 docs = [

src/scvi/external/mrvi/_model.py

@@ -10,7 +10,7 @@
 import xarray as xr
 from tqdm import tqdm
 
-from scvi import REGISTRY_KEYS
+from scvi import REGISTRY_KEYS, settings
 from scvi.data import AnnDataManager, fields
 from scvi.external.mrvi._module import MRVAE
 from scvi.external.mrvi._types import MRVIReduction
@@ -248,6 +248,14 @@ def train(
         train_kwargs["plan_kwargs"] = dict(
             deepcopy(DEFAULT_TRAIN_KWARGS["plan_kwargs"]), **plan_kwargs
         )
+        from packaging import version
+
+        if version.parse(jax.__version__) > version.parse("0.4.35"):
+            warnings.warn(
+                "Running mrVI with Jax version larger 0.4.35 can cause performance issues",
+                UserWarning,
+                stacklevel=settings.warnings_stacklevel,
+            )
         super().train(**train_kwargs)
 
     def get_latent_representation(