feat: Adding support for Annbatch

CHANGELOG.md

@@ -23,6 +23,9 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 #### Added
 
+- Add {class}`~scvi.dataloaders.AnnbatchDataModule` for out-of-core dataloading via `annbatch`,
+    enabling memory-efficient training on large-scale datasets stored as sharded Zarr collections,
+    with support for batch covariates, {pr}`3620`.
 - Add support for Python 3.14, {pr}`3563`.
 - Add support for Pandas3, {pr}`3638`.
 - Add {class}`scvi.external.DIAGVI` for integrating unpaired single-cell datasets, {pr}`3575`.

pyproject.toml

@@ -59,8 +59,8 @@ 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]", "mlx[cuda]"]
-cuda13 = ["torchvision", "torchaudio", "jax[cuda13]", "mlx[cuda13]"]
+cuda = ["torchvision", "torchaudio", "jax[cuda]", "mlx[cuda]","cupy-cuda12x"]
+cuda13 = ["torchvision", "torchaudio", "jax[cuda13]", "mlx[cuda13]","cupy-cuda13x"]
 tpu = ["torch_xla[tpu]"]
 metal = ["torchvision", "torchaudio", "jax-metal","mlx-metal"]
 
@@ -95,7 +95,7 @@ interpretability = ["captum", "shap", "decoupler"]
 # for jax support
 jax = ["jax<0.10.0", "jaxlib", "optax", "numpyro", "flax"] #TODO: unpin once it is stable again
 # for custom dataloders
-dataloaders = ["lamindb>=1.12.1", "cellxgene-census", "tiledbsoma", "tiledbsoma_ml", "torchdata"]
+dataloaders = ["lamindb>=1.12.1", "cellxgene-census", "tiledbsoma", "tiledbsoma_ml", "torchdata","annbatch[zarrs]"]
 # for diagvi
 diagvi = ["torch_geometric", "geomloss"]
 # for mlflow

src/scvi/dataloaders/__init__.py

@@ -4,7 +4,11 @@
 from ._ann_dataloader import AnnDataLoader
 from ._anncollection import CollectionAdapter
 from ._concat_dataloader import ConcatDataLoader
-from ._custom_dataloaders import MappedCollectionDataModule, TileDBDataModule
+from ._custom_dataloaders import (
+    AnnbatchDataModule,
+    MappedCollectionDataModule,
+    TileDBDataModule,
+)
 from ._data_splitting import (
     DataSplitter,
     DeviceBackedDataSplitter,
@@ -23,6 +27,7 @@
     "DataSplitter",
     "SemiSupervisedDataSplitter",
     "BatchDistributedSampler",
+    "AnnbatchDataModule",
     "MappedCollectionDataModule",
     "TileDBDataModule",
 ]