perf: remove per-minibatch overhead from the data path and the inference forward

[JESUSROYETH]

Four places paid Python work once per minibatch (or once per call) that only needs to happen once, or not at all. Nothing changes in the results: the ELBO of the introduction tutorial is the same float on both sides, and the sha256 of get_latent_representation, get_normalized_expression, get_likelihood_parameters, the differential_expression DataFrame and a minibatch from AnnTorchDataset are identical (19 of 19).

Changes

where	before	after
AnnTorchDataset.data / __getitem__	covariate DataFrames sliced with .iloc on every minibatch: 129 µs of a 214 µs batch fetch	converted to NumPy once, when data is first materialised: 3 µs
auto_move_data, BaseModuleClass.device	{p.device for p in self.parameters()} on every decorated call (3 per inference minibatch), 22 µs each on the default SCVI	next(self.parameters()).device, 1 µs
scrna_raw_counts_properties	multiply() returns COO and COO.mean(axis=0) sorts every entry first: 100 ms per DE group	.tocsr() before the mean: 13 ms, bit-identical
get_likelihood_parameters	np.concatenate(dropout_list) inside the minibatch loop, quadratic in cells (18 641 cells: 1.19 s)	concatenated once after the loop, like the other two accumulators

The device set only existed to raise if a module was split across devices, and no module in scvi-tools does that (Lightning moves the pl_module as a whole); a real split still fails with PyTorch's own device-mismatch error, just later.

Measurements

Introduction tutorial as published (heart_cell_atlas_subsampled, 18 641 × 1 200, cell_source + donor and percent_mito + percent_ribo as covariates, train_size=0.8, check_val_every_n_epoch=1), base and patched run alternately in fresh processes, median of 5 runs, delta as median (min…max) over the runs:

call	NVIDIA L4, torch 2.14.0+cu130	Intel Ice Lake 8 vCPU, torch 2.14.0+cpu
training, per epoch	1.60 → 1.47 s (−8 %, −9…−5)	2.50 → 2.37 s (−6 %, −8…−2)
validation, per epoch	0.23 → 0.20 s (−13 %)	0.30 → 0.27 s (−13 %)
get_latent_representation()	0.33 → 0.22 s (−32 %)	0.33 → 0.22 s (−34 %)
get_normalized_expression()	0.66 → 0.55 s (−16 %)	0.68 → 0.53 s (−20 %)
get_likelihood_parameters()	2.93 → 0.54 s (−81 %)	2.76 → 0.50 s (−82 %)
differential_expression(), one pair	0.38 → 0.27 s (−29 %)	0.37 → 0.26 s (−29 %)
differential_expression(mode="change"), 11 groups	12.95 → 9.20 s (−29 %)	12.44 → 8.40 s (−32 %)

The training gain needs registered covariates; the other rows apply to every model that shares the decorated forward and the DE core.

Validation

New tests: test_covariates_materialized_as_numpy, tests/module/test_decorators.py, tests/module/test_base_module.py (the DataFrame and no-parameter cases fail on current main) and test_scrna_raw_counts_properties_matches_dense_reference (all six statistics against a dense NumPy reference). tests/data, tests/dataloaders, tests/module, tests/model/test_scvi.py, test_totalvi.py, test_scanvi.py, test_peakvi.py, test_differential.py, test_pyro.py and the external/ models that read .device pass on CPU; ruff check and ruff format --check are clean. CHANGELOG entries added under 1.5.2, I will add the PR number once it exists.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 89.47%. Comparing base (73b28e4) to head (86032cf).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4017   +/-   ##
=======================================
  Coverage   89.46%   89.47%           
=======================================
  Files         247      247           
  Lines       24389    24393    +4     
=======================================
+ Hits        21820    21826    +6     
+ Misses       2569     2567    -2     

Flag	Coverage Δ
integration	75.97% <100.00%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/scvi/data/_anntorchdataset.py	97.43% <100.00%> (+0.10%)	⬆️
src/scvi/model/_utils.py	88.12% <100.00%> (ø)
src/scvi/model/base/_rnamixin.py	95.76% <100.00%> (+0.01%)	⬆️
src/scvi/module/base/_base_module.py	97.18% <100.00%> (+0.46%)	⬆️
src/scvi/module/base/_decorators.py	97.22% <100.00%> (+2.77%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ori-kron-wis]

Thank you @JESUSROYETH !
I also got 10% speed improvement in the intro tutorial.

I will validate further with other tutorials, but generally I think it safe to have this addition towards next version.