From 5343f1854f122f7f6c7f74c50e9df606b180803f Mon Sep 17 00:00:00 2001 From: ori-kron-wis Date: Sun, 28 Jun 2026 12:24:52 +0300 Subject: [PATCH 1/5] mimic PR 3886 but also pyro on_load does a lightweight no-grad guide warmup, instead of model.train --- CHANGELOG.md | 5 +++-- src/scvi/external/decipher/_model.py | 9 ++++----- src/scvi/external/decipher/_trainingplan.py | 4 ++-- src/scvi/external/resolvi/_model.py | 8 +++++++- src/scvi/module/base/_base_module.py | 20 ++++++++++++++------ tests/model/test_amortizedlda.py | 2 +- tests/model/test_pyro.py | 2 +- 7 files changed, 32 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6eae1029b7..52259e6f72 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,14 +19,15 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo #### Fixed - Fix list of metrics to be recorded in {class}`scvi.autotune.AutotuneExperiment`, {pr}`3816`. +- Fix {class}`scvi.external.RESOLVI` preparing data for every load, {pr}`3888`. #### Changed - Changed {class}`scvi.external.Tangram` backend to be in Pytorch, {pr}`3786`. - Consolidate parts of the training and data loading between {class}`~scvi.external.GIMVI` and {class}`scvi.external.DIAGVI`, {pr}`3830`. -- support validation set in DestVI training and raise clear errors for unsupported validation - in RESOLVI, {pr}`3881`. +- support validation set in {class}`scvi.model.DestVI` training and raise clear errors for + unsupported validation in {class}`scvi.external.RESOLVI`, {pr}`3881`. #### Removed diff --git a/src/scvi/external/decipher/_model.py b/src/scvi/external/decipher/_model.py index ac8009bc6e..26a043ad97 100644 --- a/src/scvi/external/decipher/_model.py +++ b/src/scvi/external/decipher/_model.py @@ -10,7 +10,6 @@ from scvi.data import AnnDataManager from scvi.data.fields import LayerField from scvi.model.base import BaseModelClass, PyroSviTrainMixin -from scvi.train import PyroTrainingPlan from scvi.utils import setup_anndata_dsp from ._module import DecipherPyroModule @@ -88,7 +87,7 @@ def train( shuffle_set_split: bool = True, batch_size: int = 128, early_stopping: bool = False, - training_plan: PyroTrainingPlan | None = None, + training_plan: DecipherTrainingPlan | None = None, datasplitter_kwargs: dict | None = None, plan_kwargs: dict | None = None, **trainer_kwargs, @@ -119,12 +118,12 @@ def train( early_stopping Perform early stopping. Additional arguments can be passed in ``**trainer_kwargs``. training_plan - Training plan instance. If ``None``, a default :class:`~scvi.train.PyroTrainingPlan` - is used. + Training plan instance. If ``None``, + a default :class:`~scvi.train.DecipherTrainingPlan` is used. datasplitter_kwargs Additional keyword arguments passed into :class:`~scvi.dataloaders.DataSplitter`. plan_kwargs - Keyword arguments for :class:`~scvi.train.PyroTrainingPlan`. + Keyword arguments for :class:`~scvi.train.DecipherTrainingPlan`. **trainer_kwargs Additional keyword arguments passed to :class:`~scvi.train.Trainer`. """ diff --git a/src/scvi/external/decipher/_trainingplan.py b/src/scvi/external/decipher/_trainingplan.py index adcb9dbbbe..099dd2585a 100644 --- a/src/scvi/external/decipher/_trainingplan.py +++ b/src/scvi/external/decipher/_trainingplan.py @@ -4,10 +4,10 @@ from scvi.module.base import ( PyroBaseModuleClass, ) -from scvi.train import LowLevelPyroTrainingPlan +from scvi.train import PyroTrainingPlan -class DecipherTrainingPlan(LowLevelPyroTrainingPlan): +class DecipherTrainingPlan(PyroTrainingPlan): """Lightning module task to train the Decipher Pyro module. Parameters diff --git a/src/scvi/external/resolvi/_model.py b/src/scvi/external/resolvi/_model.py index d68ebe65ec..06f7b07d3c 100644 --- a/src/scvi/external/resolvi/_model.py +++ b/src/scvi/external/resolvi/_model.py @@ -355,7 +355,13 @@ def setup_anndata( if prepare_data: if prepare_data_kwargs is None: prepare_data_kwargs = {} - RESOLVI._prepare_data(adata, batch_key=batch_key, **prepare_data_kwargs) + neighbors_valid = ( + "index_neighbor" in adata.obsm + and "distance_neighbor" in adata.obsm + and int(adata.obsm["index_neighbor"].max()) < adata.n_obs + ) + if not neighbors_valid: + RESOLVI._prepare_data(adata, batch_key=batch_key, **prepare_data_kwargs) anndata_fields = [ LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True), diff --git a/src/scvi/module/base/_base_module.py b/src/scvi/module/base/_base_module.py index 69b2862e31..6bc2831aad 100644 --- a/src/scvi/module/base/_base_module.py +++ b/src/scvi/module/base/_base_module.py @@ -375,17 +375,25 @@ def list_obs_plate_vars(self): return {"name": "", "in": [], "sites": {}} def on_load(self, model, **kwargs): - """Callback function run in :meth:`~scvi.model.base.BaseModelClass.load`. - - For some Pyro modules with AutoGuides, run one training step prior to loading state dict. - """ + """Callback function run in :meth:`~scvi.model.base.BaseModelClass.load`.""" pyro.clear_param_store() old_history = model.history_.copy() if model.history_ is not None else None - model.train(max_steps=1, **self.on_load_kwargs) model.history_ = old_history - if "pyro_param_store" in kwargs: + if "pyro_param_store" in kwargs and kwargs["pyro_param_store"] is not None: # For scArches shapes are changed, and we don't want to overwrite these changed shapes. pyro.get_param_store().set_state(kwargs["pyro_param_store"]) + # AutoGuide parameters are lazy — they only exist as nn.Module attributes after the + # first forward call. Run one no-grad guide pass so load_state_dict finds the keys. + if model.adata is not None: + if hasattr(self.model, "n_obs"): + self.model.n_obs = model.adata.n_obs + if hasattr(self.guide, "n_obs"): + self.guide.n_obs = model.adata.n_obs + dl = model._make_data_loader(model.adata, batch_size=2) + batch = next(iter(dl)) + args, guide_kwargs = self._get_fn_args_from_batch(batch) + with torch.no_grad(): + self.guide(*args, **guide_kwargs) def create_predictive( self, diff --git a/tests/model/test_amortizedlda.py b/tests/model/test_amortizedlda.py index 431f843bda..0fb87e23b6 100644 --- a/tests/model/test_amortizedlda.py +++ b/tests/model/test_amortizedlda.py @@ -96,7 +96,7 @@ def test_lda_model(n_topics: int): @pytest.mark.parametrize("n_topics", [5]) -def test_lda_model_save_load(save_path: str, n_topics: int): +def test_lda_model_save_load(n_topics: int, save_path: str): adata = synthetic_iid() AmortizedLDA.setup_anndata(adata) mod = AmortizedLDA(adata, n_topics=n_topics) diff --git a/tests/model/test_pyro.py b/tests/model/test_pyro.py index 1d0cacf48c..e17787d340 100644 --- a/tests/model/test_pyro.py +++ b/tests/model/test_pyro.py @@ -339,7 +339,7 @@ def test_pyro_bayesian_save_load(save_path): linear_median = quants["linear.weight"][0].detach().cpu().numpy() model_save_path = os.path.join(save_path, "test_pyro_bayesian/") - mod.save(model_save_path) + mod.save(model_save_path, overwrite=True, save_anndata=True) # Test setting `on_load_kwargs` mod.module.on_load_kwargs = {"batch_size": 8} From 81d9261c616eb2fdc996dba52766b971fdc481d1 Mon Sep 17 00:00:00 2001 From: ori-kron-wis Date: Sun, 28 Jun 2026 13:06:29 +0300 Subject: [PATCH 2/5] fix: replace on_load train(max_steps=1) with no-grad guide warmup and guard RESOLVI _prepare_data against stale-config cache --- src/scvi/external/resolvi/_model.py | 10 +++++++++- src/scvi/module/base/_base_module.py | 20 ++++++++++++++++---- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/scvi/external/resolvi/_model.py b/src/scvi/external/resolvi/_model.py index 06f7b07d3c..e3eea946d2 100644 --- a/src/scvi/external/resolvi/_model.py +++ b/src/scvi/external/resolvi/_model.py @@ -355,13 +355,21 @@ def setup_anndata( if prepare_data: if prepare_data_kwargs is None: prepare_data_kwargs = {} + effective_config = {"batch_key": batch_key, **prepare_data_kwargs} + stored_config = adata.uns.get("_resolvi_prepare_data_config", None) neighbors_valid = ( - "index_neighbor" in adata.obsm + stored_config == effective_config + and "index_neighbor" in adata.obsm and "distance_neighbor" in adata.obsm and int(adata.obsm["index_neighbor"].max()) < adata.n_obs ) if not neighbors_valid: + print( + "Preparing data for training. This may take a while. " + "RAPIDS SingleCell will be used if installed." + ) RESOLVI._prepare_data(adata, batch_key=batch_key, **prepare_data_kwargs) + adata.uns["_resolvi_prepare_data_config"] = effective_config anndata_fields = [ LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True), diff --git a/src/scvi/module/base/_base_module.py b/src/scvi/module/base/_base_module.py index 6bc2831aad..2426ccacc8 100644 --- a/src/scvi/module/base/_base_module.py +++ b/src/scvi/module/base/_base_module.py @@ -380,10 +380,22 @@ def on_load(self, model, **kwargs): old_history = model.history_.copy() if model.history_ is not None else None model.history_ = old_history if "pyro_param_store" in kwargs and kwargs["pyro_param_store"] is not None: - # For scArches shapes are changed, and we don't want to overwrite these changed shapes. - pyro.get_param_store().set_state(kwargs["pyro_param_store"]) - # AutoGuide parameters are lazy — they only exist as nn.Module attributes after the - # first forward call. Run one no-grad guide pass so load_state_dict finds the keys. + store_state = kwargs["pyro_param_store"] + # Remap tensors to CPU so the warmup below runs consistently on a CPU + # module/batch. load_state_dict + model.to_device() handle final placement. + # (For scArches the saved store has old shapes — those are intentionally + # restored here; the query model's new shapes are set before load_state_dict.) + cpu_state = { + "params": { + k: v.cpu() if isinstance(v, torch.Tensor) else v + for k, v in store_state.get("params", {}).items() + }, + "constraints": store_state.get("constraints", {}), + } + pyro.get_param_store().set_state(cpu_state) + # AutoGuide parameters are lazy — they only exist as nn.Module attributes + # after the first forward call. Run one no-grad guide pass so + # load_state_dict finds the expected keys. if model.adata is not None: if hasattr(self.model, "n_obs"): self.model.n_obs = model.adata.n_obs From 1a0c68a0a6914d7a680bbd0dfbc5723347297219 Mon Sep 17 00:00:00 2001 From: ori-kron-wis Date: Sun, 28 Jun 2026 13:17:48 +0300 Subject: [PATCH 3/5] =?UTF-8?q?warmup=20against=20empty=20store=20?= =?UTF-8?q?=E2=86=92=20restore=20saved=20state=20=E2=86=92=20load=5Fstate?= =?UTF-8?q?=5Fdict=20+=20to=5Fdevice.=20works=20for=20CPU=20or=20GPU?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/scvi/module/base/_base_module.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/src/scvi/module/base/_base_module.py b/src/scvi/module/base/_base_module.py index 2426ccacc8..ffad30eafd 100644 --- a/src/scvi/module/base/_base_module.py +++ b/src/scvi/module/base/_base_module.py @@ -381,21 +381,11 @@ def on_load(self, model, **kwargs): model.history_ = old_history if "pyro_param_store" in kwargs and kwargs["pyro_param_store"] is not None: store_state = kwargs["pyro_param_store"] - # Remap tensors to CPU so the warmup below runs consistently on a CPU - # module/batch. load_state_dict + model.to_device() handle final placement. - # (For scArches the saved store has old shapes — those are intentionally - # restored here; the query model's new shapes are set before load_state_dict.) - cpu_state = { - "params": { - k: v.cpu() if isinstance(v, torch.Tensor) else v - for k, v in store_state.get("params", {}).items() - }, - "constraints": store_state.get("constraints", {}), - } - pyro.get_param_store().set_state(cpu_state) # AutoGuide parameters are lazy — they only exist as nn.Module attributes - # after the first forward call. Run one no-grad guide pass so - # load_state_dict finds the expected keys. + # after the first forward call. Run one no-grad guide pass BEFORE restoring + # the saved store so the CPU module always sees a device-neutral empty store. + # Restoring first would expose CUDA constraint tensors (e.g. AffineTransform + # bounds derived from CUDA buffers) that clash with the CPU batch tensors. if model.adata is not None: if hasattr(self.model, "n_obs"): self.model.n_obs = model.adata.n_obs @@ -406,6 +396,10 @@ def on_load(self, model, **kwargs): args, guide_kwargs = self._get_fn_args_from_batch(batch) with torch.no_grad(): self.guide(*args, **guide_kwargs) + # Restore the saved store after warmup. For scArches the saved store + # carries old reference shapes — those are intentionally kept here. + # load_state_dict + model.to_device() handle the final device placement. + pyro.get_param_store().set_state(store_state) def create_predictive( self, From 447f7815a2e180b5ed96e4564574874287516b08 Mon Sep 17 00:00:00 2001 From: ori-kron-wis Date: Sun, 28 Jun 2026 13:20:54 +0300 Subject: [PATCH 4/5] adding decipher save/load test --- tests/external/decipher/test_decipher.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/external/decipher/test_decipher.py b/tests/external/decipher/test_decipher.py index d95e9df9e2..6dc6b0d64c 100644 --- a/tests/external/decipher/test_decipher.py +++ b/tests/external/decipher/test_decipher.py @@ -68,6 +68,27 @@ def test_decipher_model_methods(adata): assert z_cov.shape == (adata.n_vars, model.module.dim_z) +def test_decipher_save_load(adata, save_path): + import os + + Decipher.setup_anndata(adata) + model = Decipher(adata) + model.train(max_epochs=2, train_size=0.5) + + v_before = model.get_latent_representation(give_z=False) + z_before = model.get_latent_representation(give_z=True) + + path = os.path.join(save_path, "test_decipher") + model.save(path, overwrite=True, save_anndata=True) + model2 = Decipher.load(path) + + v_after = model2.get_latent_representation(give_z=False) + z_after = model2.get_latent_representation(give_z=True) + + np.testing.assert_array_almost_equal(v_before, v_after, decimal=4) + np.testing.assert_array_almost_equal(z_before, z_after, decimal=4) + + def test_decipher_trajectory(adata): Decipher.setup_anndata(adata) model = Decipher(adata) From 08889831a32a505bbaf9a8f4e01bed7583fd02e9 Mon Sep 17 00:00:00 2001 From: ori-kron-wis Date: Sun, 28 Jun 2026 13:44:18 +0300 Subject: [PATCH 5/5] fix changlog for annbatch --- CHANGELOG.md | 4 ++-- docs/api/developer.md | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 52259e6f72..12291316b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,12 +14,12 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo enabling memory-efficient training on large-scale datasets stored as sharded Zarr collections, with support for batch covariates, {pr}`3620`. - Add support for rapids-singlecell, {pr}`3811`. -- Add {class}`scvi.external.CytoVI` KNN imputation backend option to be cuML, {pr}`3821`. +- Add {class}`scvi.external.CYTOVI` KNN imputation backend option to be cuML, {pr}`3821`. #### Fixed - Fix list of metrics to be recorded in {class}`scvi.autotune.AutotuneExperiment`, {pr}`3816`. -- Fix {class}`scvi.external.RESOLVI` preparing data for every load, {pr}`3888`. +- Fix {class}`scvi.external.RESOLVI` preparing data for every load, {pr}`3887`. #### Changed diff --git a/docs/api/developer.md b/docs/api/developer.md index 3a33042a2a..510535beb4 100644 --- a/docs/api/developer.md +++ b/docs/api/developer.md @@ -72,6 +72,7 @@ DataLoaders for loading tensors from AnnData objects. DataSplitters for splittin dataloaders.AnnDataLoader dataloaders.AnnTorchDataset + dataloaders.AnnbatchDataModule dataloaders.CollectionAdapter dataloaders.ConcatDataLoader dataloaders.DataSplitter