fix: reject zero multi-GPU workers

docs/release-notes/0.17.0.md

@@ -2,6 +2,7 @@
 
 ```{rubric} Bug fixes
 ```
+* Fix {func}`~rapids_singlecell.tl.score_genes` with a custom `gene_pool` on CSR and Dask inputs, and its per-gene means on Dask dense inputs that contain NaNs {pr}`787` {smaller}`S Dicks`
 * Make the ``lanczos`` SVD breakdown check scale-aware; the previous fixed threshold never triggered in float32. {pr}`755` {smaller}`S Dicks`
 * Fix {func}`~rapids_singlecell.pp.harmony_integrate` with RMM managed memory by making multi-key clustering arrays optional and allocating them only when needed. {pr}`766` {smaller}`A Holly & S Dicks`
 * Fix {func}`~rapids_singlecell.tl.umap` on neighbor graphs whose ``metric`` ``cuml`` does not know {pr}`768` {smaller}`S Dicks`

src/rapids_singlecell/_utils/_multi_gpu.py

@@ -188,6 +188,9 @@ def _split_pairs(
     list
         List of (pair_left, pair_right) tuples for each device
     """
+    if n_devices < 1:
+        raise ValueError("n_devices must be at least 1")
+
     n_pairs = len(pair_left)
 
     if n_pairs == 0:

src/rapids_singlecell/tools/_utils.py

@@ -205,7 +205,7 @@ def _nan_mean(X, axis=0, *, mask=None, n_features=None):
                     mask = cp.ones(X.shape[1], dtype=cp.bool_)
                 mean = _nan_mean_minor(
                     X, major, minor, mask=mask, n_features=n_features
-                )
+                )[mask]
             elif isspmatrix_csc(X):
                 if mask is not None:
                     X = X[:, mask]
@@ -245,7 +245,7 @@ def _nan_mean(X, axis=0, *, mask=None, n_features=None):
                 n_features = major
                 mean = _nan_mean_minor_dask_sparse(
                     X, major, minor, mask=mask, n_features=n_features
-                )
+                )[mask]
             elif axis == 1:
                 n_features = minor if n_features is None else n_features
                 mean = _nan_mean_major_dask_sparse(
@@ -256,7 +256,7 @@ def _nan_mean(X, axis=0, *, mask=None, n_features=None):
         elif isinstance(X._meta, cp.ndarray):
             if mask is None:
                 mask = cp.ones(X.shape[1], dtype=cp.bool_)
-            if n_features is None:
+            if axis == 0 or n_features is None:
                 n_features = X.shape[axis]
             mean = _nan_mean_dense_dask(X, axis, mask=mask, n_features=n_features)
             # raise NotImplementedError("Dask dense arrays are not supported yet")

tests/test_multi_gpu_utils.py

@@ -23,6 +23,13 @@ def test_empty_pairs(self):
             assert len(left) == 0
             assert len(right) == 0
 
+    def test_zero_devices_is_rejected(self):
+        pair_left = cp.array([], dtype=cp.int32)
+        pair_right = cp.array([], dtype=cp.int32)
+
+        with pytest.raises(ValueError, match="n_devices must be at least 1"):
+            _split_pairs(pair_left, pair_right, n_devices=0)
+
     def test_single_device_returns_all_pairs(self):
         """Single device gets all pairs."""
         pair_left = cp.array([0, 0, 1], dtype=cp.int32)