diff --git a/CHANGELOG.md b/CHANGELOG.md index 37ad4537d3..303cbbbde4 100644 --- a/CHANGELOG.md +++ b/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`, diff --git a/pyproject.toml b/pyproject.toml index 2cbd124d93..ab45ee9dd4 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/scvi/data/_utils.py b/src/scvi/data/_utils.py index 1b5c5a5253..0767d4f277 100644 --- a/src/scvi/data/_utils.py +++ b/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, diff --git a/src/scvi/dataloaders/_data_splitting.py b/src/scvi/dataloaders/_data_splitting.py index 3abdaf3518..36ae6c6ad3 100644 --- a/src/scvi/dataloaders/_data_splitting.py +++ b/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 diff --git a/src/scvi/module/_vae.py b/src/scvi/module/_vae.py index 057bfd680c..4777988348 100644 --- a/src/scvi/module/_vae.py +++ b/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) diff --git a/src/scvi/nn/_base_components.py b/src/scvi/nn/_base_components.py index a4ef76b78c..939b7a2b1c 100644 --- a/src/scvi/nn/_base_components.py +++ b/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): diff --git a/tests/dataloaders/sparse_utils.py b/tests/dataloaders/sparse_utils.py index 67f4aecf14..d39ab689cb 100644 --- a/tests/dataloaders/sparse_utils.py +++ b/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, diff --git a/tests/dataloaders/test_datasplitter.py b/tests/dataloaders/test_datasplitter.py index 0115362b03..67655492af 100644 --- a/tests/dataloaders/test_datasplitter.py +++ b/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") diff --git a/tests/model/test_scvi.py b/tests/model/test_scvi.py index 107746d983..a6e027a591 100644 --- a/tests/model/test_scvi.py +++ b/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 diff --git a/tests/module/test_vae.py b/tests/module/test_vae.py index 3ba2a8a86d..e79bc686d3 100644 --- a/tests/module/test_vae.py +++ b/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)