feat: supporting sparse mode in DataSplitter

CHANGELOG.md

@@ -9,6 +9,12 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 #### Added
 
+- Add `sparse_mode` to {class}`~scvi.dataloaders.DataSplitter` (`"OFF"`, `"TRANSPORT"`,
+    `"INPUT_CSR"`, `"AUTO"`). When `load_sparse_tensor=True`, `"INPUT_CSR"` keeps the count
+    matrix as a sparse CSR tensor through `log1p` and the first encoder linear layer
+    ({class}`~scvi.nn.FCLayers`) before densifying, and `"AUTO"` selects `"INPUT_CSR"` or
+    `"TRANSPORT"` from the measured per-batch density. Default `"TRANSPORT"` preserves prior
+    behavior, {pr}`3840`.
 - Add [scvi-tools MCP](https://scvi-tools-mcp.readthedocs.io/en/latest/index.html) package
     that gives any MCP-compatible LLM access to scvi-tools knowledge.
 - Add {class}`~scvi.dataloaders.AnnbatchDataModule` for out-of-core dataloading via `annbatch`,

pyproject.toml

@@ -134,6 +134,7 @@ omit = [
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
+pythonpath = ["."]
 xfail_strict = true
 markers = [
     "internet: mark tests that requires internet access",

src/scvi/data/_utils.py

@@ -9,6 +9,7 @@
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp_sparse
+import torch
 from anndata import AnnData
 from anndata.abc import CSCDataset, CSRDataset
 from anndata.io import read_elem
@@ -87,6 +88,18 @@ def scipy_to_torch_sparse(x: sp_sparse.csr_matrix | sp_sparse.csc_matrix) -> Ten
         )
 
 
+def normalize_to_csr(x: Tensor) -> Tensor:
+    """Return a sparse CSR tensor, converting from CSC if needed.
+
+    CSR is the only sparse layout with reliable autograd support for
+    :func:`torch.sparse.mm` across supported PyTorch versions (see issue #2550).
+    CSR input and dense input are returned unchanged.
+    """
+    if x.layout is torch.sparse_csc:
+        return x.to_sparse_csr()
+    return x
+
+
 def get_anndata_attribute(
     adata: AnnOrMuData,
     attr_name: str,

src/scvi/dataloaders/_data_splitting.py

@@ -14,12 +14,19 @@
 
 from scvi import REGISTRY_KEYS, settings
 from scvi.data import AnnDataManager
-from scvi.data._utils import get_anndata_attribute
+from scvi.data._utils import get_anndata_attribute, normalize_to_csr
 from scvi.dataloaders._ann_dataloader import AnnDataLoader
 from scvi.dataloaders._semi_dataloader import SemiSupervisedDataLoader
 from scvi.model._utils import parse_device_args
 from scvi.utils._docstrings import devices_dsp
 
+# Sparse handling modes for DataSplitter.on_after_batch_transfer.
+#   OFF/TRANSPORT : densify all sparse tensors after transfer (legacy behavior).
+#   INPUT_CSR     : keep REGISTRY_KEYS.X as CSR (CSC normalized to CSR), densify the rest.
+#   AUTO          : INPUT_CSR if measured X density <= SPARSE_AUTO_DENSITY, else TRANSPORT.
+SPARSE_MODES = ("OFF", "TRANSPORT", "INPUT_CSR", "AUTO")
+SPARSE_AUTO_DENSITY = 0.10  # crossover from benchmark_sparse.py on CUDA
+
 
 def validate_data_split(
     n_samples: int,
@@ -231,6 +238,7 @@ def __init__(
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,
+        sparse_mode: str = "TRANSPORT",
         pin_memory: bool = False,
         external_indexing: list[np.array, np.array, np.array] | None = None,
         **kwargs,
@@ -242,6 +250,9 @@ def __init__(
         self.validation_size = validation_size
         self.shuffle_set_split = shuffle_set_split
         self.load_sparse_tensor = load_sparse_tensor
+        if sparse_mode not in SPARSE_MODES:
+            raise ValueError(f"`sparse_mode` must be one of {SPARSE_MODES}, got {sparse_mode!r}.")
+        self.sparse_mode = sparse_mode
         self.drop_last = kwargs.pop("drop_last", False)
         self.data_loader_kwargs = kwargs
         self.pin_memory = pin_memory
@@ -329,12 +340,30 @@ def test_dataloader(self):
             pass
 
     def on_after_batch_transfer(self, batch, dataloader_idx):
-        """Converts sparse tensors to dense if necessary."""
-        if self.load_sparse_tensor:
-            for key, val in batch.items():
-                layout = val.layout if isinstance(val, torch.Tensor) else None
-                if layout is torch.sparse_csr or layout is torch.sparse_csc:
-                    batch[key] = val.to_dense()
+        """Converts sparse tensors to dense, honoring ``sparse_mode``.
+
+        In ``INPUT_CSR`` (or ``AUTO`` below the density threshold) mode the
+        :attr:`~scvi.REGISTRY_KEYS.X_KEY` tensor is kept as a sparse CSR tensor
+        (converting from CSC if needed) so it can flow through ``log1p`` and the
+        first encoder linear layer; all other tensors are densified as before.
+        """
+        if not self.load_sparse_tensor:
+            return batch
+
+        for key, val in batch.items():
+            layout = val.layout if isinstance(val, torch.Tensor) else None
+            if layout is not torch.sparse_csr and layout is not torch.sparse_csc:
+                continue
+
+            keep_sparse = False
+            if self.sparse_mode in ("INPUT_CSR", "AUTO") and key == REGISTRY_KEYS.X_KEY:
+                if self.sparse_mode == "AUTO":
+                    density = val._nnz() / (val.shape[0] * val.shape[1])
+                    keep_sparse = density <= SPARSE_AUTO_DENSITY
+                else:
+                    keep_sparse = True
+
+            batch[key] = normalize_to_csr(val) if keep_sparse else val.to_dense()
 
         return batch
 

src/scvi/module/_vae.py

@@ -369,11 +369,18 @@ def _regular_inference(
         """Run the regular inference process."""
         x_ = x
         if self.use_observed_lib_size:
-            library = torch.log(x.sum(1)).unsqueeze(1)
+            # keepdim avoids the unsupported keepdim=False reduction on sparse CSR
+            # inputs (INPUT_CSR mode); identical to x.sum(1).unsqueeze(1) for dense x.
+            x_sum = x.sum(dim=1, keepdim=True)
+            if x_sum.layout is not torch.strided:
+                x_sum = x_sum.to_dense()
+            library = torch.log(x_sum)
         if self.log_variational:
             x_ = torch.log1p(x_)
 
         if cont_covs is not None and self.encode_covariates:
+            if x_.layout is not torch.strided:
+                x_ = x_.to_dense()
             encoder_input = torch.cat((x_, cont_covs), dim=-1)
         else:
             encoder_input = x_
@@ -384,6 +391,8 @@ def _regular_inference(
 
         if self.batch_representation == "embedding" and self.encode_covariates:
             batch_rep = self.compute_embedding(REGISTRY_KEYS.BATCH_KEY, batch_index)
+            if encoder_input.layout is not torch.strided:
+                encoder_input = encoder_input.to_dense()
             encoder_input = torch.cat([encoder_input, batch_rep], dim=-1)
             qz, z = self.z_encoder(encoder_input, *categorical_input)
         else:
@@ -560,6 +569,10 @@ def loss(
         from torch.distributions import kl_divergence
 
         x = tensors[REGISTRY_KEYS.X_KEY]
+        if x.layout is not torch.strided:
+            # INPUT_CSR mode keeps x sparse for the encoder; the ZINB reconstruction
+            # target requires dense x (sparse reconstruction loss is future work, #1038).
+            x = x.to_dense()
         kl_divergence_z = kl_divergence(
             inference_outputs[MODULE_KEYS.QZ_KEY], generative_outputs[MODULE_KEYS.PZ_KEY]
         ).sum(dim=-1)

src/scvi/nn/_base_components.py

@@ -178,7 +178,35 @@ def forward(self, x: torch.Tensor, *cat_list: int, cont: torch.Tensor | None = N
                     one_hot_cat = cat  # cat has already been one_hot encoded
                 one_hot_cat_list += [one_hot_cat]
         cov_list = cont_list + one_hot_cat_list
+
+        # Sparse fast path: keep a CSR input sparse through the first linear layer
+        # (x @ W.T), densifying only afterwards for BatchNorm/activation/dropout.
+        sparse_input = isinstance(x, torch.Tensor) and x.layout in (
+            torch.sparse_csr,
+            torch.sparse_csc,
+        )
+        if sparse_input and x.layout is torch.sparse_csc:
+            # CSC@dense backward is unreliable on older PyTorch (issue #2550).
+            x = x.to_sparse_csr()
+
         for i, layers in enumerate(self.fc_layers):
+            if i == 0 and sparse_input:
+                linear = layers[0]
+                weight = linear.weight  # [n_out, n_in (+ covariates)]
+                n_in = x.shape[1]
+                # Split the linear: sparse x-block plus dense covariate-block. This
+                # reproduces the dense path's concat-then-matmul without densifying x.
+                h = torch.sparse.mm(x, weight[:, :n_in].t())
+                if cov_list and self.inject_into_layer(i):
+                    cov = torch.cat(cov_list, dim=-1).to(dtype=weight.dtype)
+                    h = h + cov @ weight[:, n_in:].t()
+                if linear.bias is not None:
+                    h = h + linear.bias
+                x = h  # dense from here on
+                for layer in layers[1:]:
+                    if layer is not None:
+                        x = layer(x)
+                continue
             for layer in layers:
                 if layer is not None:
                     if isinstance(layer, nn.BatchNorm1d):

tests/dataloaders/sparse_utils.py

@@ -35,6 +35,23 @@ def on_after_batch_transfer(self, batch, dataloader_idx):
         return batch
 
 
+class TestInputCSRDataSplitter(scvi.dataloaders.DataSplitter):
+    """Asserts X stays sparse CSR after transfer (INPUT_CSR mode); rest is dense."""
+
+    def on_after_batch_transfer(self, batch, dataloader_idx):
+        batch = super().on_after_batch_transfer(batch, dataloader_idx)
+
+        X = batch.get(scvi.REGISTRY_KEYS.X_KEY)
+        assert isinstance(X, torch.Tensor)
+        assert X.layout is torch.sparse_csr  # kept sparse; CSC normalized to CSR
+
+        for key, val in batch.items():
+            if key != scvi.REGISTRY_KEYS.X_KEY and isinstance(val, torch.Tensor):
+                assert val.layout is torch.strided
+
+        return batch
+
+
 class TestSparseTrainingPlan(scvi.train.TrainingPlan):
     def training_step(self, batch, batch_idx):
         pass
@@ -94,13 +111,22 @@ def train(
         devices: int | list[int] | str = "auto",
         expected_sparse_layout: Literal["csr", "csc"] = None,
         external_indexing: list[np.array, np.array, np.array] | None = None,
+        sparse_mode: str = "TRANSPORT",
     ):
-        data_splitter = TestSparseDataSplitter(
-            self.adata_manager,
-            expected_sparse_layout=expected_sparse_layout,
-            load_sparse_tensor=True,
-            external_indexing=external_indexing,
-        )
+        if sparse_mode == "INPUT_CSR":
+            data_splitter = TestInputCSRDataSplitter(
+                self.adata_manager,
+                load_sparse_tensor=True,
+                sparse_mode="INPUT_CSR",
+                external_indexing=external_indexing,
+            )
+        else:
+            data_splitter = TestSparseDataSplitter(
+                self.adata_manager,
+                expected_sparse_layout=expected_sparse_layout,
+                load_sparse_tensor=True,
+                external_indexing=external_indexing,
+            )
         training_plan = TestSparseTrainingPlan(self.module)
         runner = TrainRunner(
             self,

tests/dataloaders/test_datasplitter.py

@@ -228,3 +228,24 @@ def test_datasplitter_load_sparse_tensor(
         expected_sparse_layout=sparse_format.split("_")[0],
         external_indexing=[np.array(train_ind), np.array(valid_ind)],
     )
+
+
+@pytest.mark.parametrize("sparse_format", ["csr_matrix", "csc_matrix"])
+def test_datasplitter_input_csr_mode(
+    sparse_format: str,
+    accelerator: str,
+    devices: list | str | int,
+):
+    # INPUT_CSR keeps X sparse (CSR) past on_after_batch_transfer; CSC is
+    # normalized to CSR. The helper splitter asserts the post-transfer layout.
+    adata = synthetic_iid(sparse_format=sparse_format)
+    TestSparseModel.setup_anndata(adata)
+    model = TestSparseModel(adata)
+    model.train(accelerator=accelerator, devices=devices, sparse_mode="INPUT_CSR")
+
+
+def test_datasplitter_invalid_sparse_mode():
+    adata = synthetic_iid()
+    manager = generic_setup_adata_manager(adata)
+    with pytest.raises(ValueError, match="sparse_mode"):
+        DataSplitter(manager, sparse_mode="NOT_A_MODE")

tests/model/test_scvi.py

@@ -1711,3 +1711,54 @@ def test_scvi_mlflow(
     mlflow_log_artifact(
         model_path + "/" + run_name + "_" + SAVE_KEYS.MODEL_FNAME, run_id=model.run_id
     )
+
+
+@pytest.mark.parametrize(
+    "model_kwargs",
+    [
+        {},
+        {"encode_covariates": True},  # covers sparse one-hot covariate matmul dtype cast
+        {"encode_covariates": True, "use_cont": True},  # cont covs concat with sparse x_
+        {"encode_covariates": True, "batch_representation": "embedding"},  # embedding concat
+    ],
+)
+def test_scvi_input_csr_equivalence(
+    accelerator: str,
+    devices: list | str | int,
+    model_kwargs: dict,
+):
+    # INPUT_CSR sparse path (encoder first layer) must match the dense path
+    # numerically across covariate configurations (regressions for PR #3840
+    # Codex review: covariate dtype cast and pre-encoder concat with sparse x).
+    use_cont = model_kwargs.pop("use_cont", False)
+
+    def make_adata():
+        adata = synthetic_iid(sparse_format="csr_matrix")
+        if use_cont:
+            adata.obs["cont1"] = np.random.RandomState(0).normal(size=adata.n_obs)
+        return adata
+
+    def run(sparse_mode):
+        scvi.settings.seed = 0
+        adata = make_adata()
+        SCVI.setup_anndata(
+            adata,
+            batch_key="batch",
+            continuous_covariate_keys=["cont1"] if use_cont else None,
+        )
+        model = SCVI(adata, n_latent=5, **model_kwargs)
+        train_kwargs = {
+            "max_epochs": 3,
+            "accelerator": accelerator,
+            "devices": devices,
+            "enable_progress_bar": False,
+        }
+        if sparse_mode is not None:
+            train_kwargs["load_sparse_tensor"] = True
+            train_kwargs["datasplitter_kwargs"] = {"sparse_mode": sparse_mode}
+        model.train(**train_kwargs)
+        return float(model.history["elbo_train"].iloc[-1].item())
+
+    dense = run(None)
+    input_csr = run("INPUT_CSR")
+    assert abs(dense - input_csr) / abs(dense) < 0.01

tests/module/test_vae.py

@@ -3,6 +3,7 @@
 
 from scvi import REGISTRY_KEYS
 from scvi.module import VAE
+from scvi.module._constants import MODULE_KEYS
 
 
 @pytest.mark.parametrize("n_samples", [1, 2, 3])
@@ -31,3 +32,43 @@ def test_sample(
         assert x_hat.shape == (batch_size, n_input, n_samples)
     else:
         assert x_hat.shape == (batch_size, n_input)
+
+
+@pytest.mark.parametrize("pre_encoder_covariate", ["continuous", "embedding_batch"])
+def test_regular_inference_accepts_sparse_x_with_pre_encoder_covariates(
+    pre_encoder_covariate: str,
+):
+    n_input = 5
+    batch_size = 4
+    use_continuous_covariate = pre_encoder_covariate == "continuous"
+    vae = VAE(
+        n_input=n_input,
+        n_batch=2,
+        n_hidden=8,
+        n_latent=3,
+        n_continuous_cov=int(use_continuous_covariate),
+        encode_covariates=True,
+        batch_representation=(
+            "embedding" if pre_encoder_covariate == "embedding_batch" else "one-hot"
+        ),
+        use_batch_norm="none",
+        use_layer_norm="none",
+    )
+    x = torch.tensor(
+        [
+            [1.0, 0.0, 2.0, 0.0, 1.0],
+            [0.0, 3.0, 0.0, 1.0, 0.0],
+            [2.0, 0.0, 1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0, 2.0, 1.0],
+        ]
+    )
+    batch_index = torch.tensor([[0], [1], [0], [1]])
+    cont_covs = torch.randn(batch_size, 1) if use_continuous_covariate else None
+
+    inference_outputs = vae._regular_inference(
+        x.to_sparse_csr(),
+        batch_index,
+        cont_covs=cont_covs,
+    )
+
+    assert inference_outputs[MODULE_KEYS.Z_KEY].shape == (batch_size, 3)