Keep BiCV's held-out cell blocks lazy for duck-typed backends - #556
Merged
Merged
Conversation
_bicv_trial evaluated a rank's fit by row-restricting X.X with a boolean
cell mask (X_mat[train_cell_mask]/X_mat[test_cell_mask]) and handing the
result to rmatmul/calc_W, deliberately avoiding materializing the raw data
("reaches the raw data through products... rather than materialising a
block of it"). That holds for a plain ndarray or scipy-sparse X_mat, but
not for a vsparse normalized view (e.g. BAL-Pf2's lazy-normalized-view
AnnData): bracket indexing such a view is documented to always eagerly
build a dense ndarray for the selection. At BAL-Pf2's real scale (1.3M
cells), a single train/test split (~50% of all cells) would materialize
on the order of tens of GB, once or twice per BiCV trial, across every
rank/repeat in a sweep -- never surfaced before because an unrelated
vsparse memory issue always killed these runs earlier in the pipeline.
Adds _restrict_rows(X_mat, mask), which uses vsparse's new
select(recalculate=False) (meyer-lab/vsparse#50) to stay a genuinely lazy
view for a duck-typed backend, falling back to ordinary indexing for
plain dense/sparse X_mat (unaffected, still cheap for those). Also adds a
third branch to _test_block_moments for the same duck-typed case: selects
the cell subset lazily, then streams it in bounded row chunks rather than
ever materializing the whole subset as one dense block.
Temporarily pins vsparse to its (as yet unmerged/unreleased)
select-without-recalculate branch in pyproject.toml -- drop once
meyer-lab/vsparse#50 merges and releases, reverting to the plain PyPI
version constraint.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Avoids a conflicting-git-ref resolution error for downstream consumers (e.g. BAL-Pf2) that pin vsparse's combined bal-pf2-gpu-testing branch (carrying both #49's matmul kernel and #50's select(recalculate=False)) rather than #50's branch alone. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
5 tasks
The matmul kernel and lazy select(recalculate=False) this branch needed have been released, so drop the temporary git pin to the combined testing branch in favor of the PyPI release. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_bicv_trialevaluates a rank's fit by row-restrictingX.Xwith a booleancell mask (
X_mat[train_cell_mask]/X_mat[test_cell_mask]) and handing theresult to
rmatmul/calc_W, deliberately to avoid materializing the rawdata (see the existing comment: "reaches the raw data through products...
rather than materialising a block of it"). That holds for a plain
ndarrayor scipy-sparse
X_mat, but not for a vsparse normalized view (e.g.BAL-Pf2's lazy-normalized-view
AnnData): bracket indexing such a view isdocumented to always eagerly build a dense
ndarrayfor the selection.At BAL-Pf2's real scale (1.3M cells), a single train/test split (~50% of all
cells) would materialize on the order of tens of GB, once or twice per BiCV
trial, across every rank/repeat in a 100-rank sweep. This never surfaced
before because an unrelated vsparse memory issue (fixed in
meyer-lab/vsparse#49) always killed these runs earlier in the pipeline,
during compression -- this only became visible once that was fixed and a
real subsample run got far enough to reach the evaluation step.
Changes
_restrict_rows(X_mat, mask): uses vsparse's newselect(recalculate=False)(Add select(recalculate=False): a lazy, stats-preserving row/col selection vsparse#50) to stay a genuinelylazy view for a duck-typed backend, falling back to ordinary bracket
indexing for plain dense/sparse
X_mat(unaffected, still cheap forthose). Used at all three
X_mat[mask]sites in_bicv_trial._test_block_momentsgains a third branch for the same duck-typed case:selects the cell subset lazily, then streams it in bounded row chunks
(
_MOMENT_CHUNK_BUDGET_BYTES) rather than ever materializing the wholesubset as one dense block -- mirroring the existing scipy-sparse branch's
own streaming design, just against a different backend's primitives.
pyproject.toml: temporarily pinsvsparseto its (as yetunmerged/unreleased)
select-without-recalculatebranch -- drop onceAdd select(recalculate=False): a lazy, stats-preserving row/col selection vsparse#50 merges and releases, reverting to the plain PyPI
version constraint.
Test plan
uv run pytest scrise/tests/-- 131 passed, 1 skipped (pre-existingskip, unrelated)
uv run ruff check ./uv run ruff format --check .uv run ty check scrise/test_rank_selection.py:_restrict_rowsmatchesbracket indexing for a normalized view while staying lazy (both VCSR
and VCSC), is a no-op passthrough for plain dense/sparse arrays,
_test_block_moments's new branch matches a dense reference, streamsacross multiple row chunks correctly, and the streaming loop's own
peak memory is bounded by the chunk size rather than the selected
block's total size (isolated from the one-time, pre-existing cost of
building the lazy selection itself, which is a separate, bounded,
single-occurrence cost this PR doesn't change).
Related
minor-axis kernel)
select(recalculate=False), which this PRdepends on)