Skip to content

Keep BiCV's held-out cell blocks lazy for duck-typed backends - #556

Merged
aarmey merged 3 commits into
mainfrom
vsparse-lazy-bicv-blocks
Sep 14, 2026
Merged

aarmey merged 3 commits into
mainfrom
vsparse-lazy-bicv-blocks

Conversation

@aarmey

@aarmey aarmey commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

_bicv_trial evaluates 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 to avoid materializing the raw
data (see the existing comment: "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 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 new
    select(recalculate=False) (Add select(recalculate=False): a lazy, stats-preserving row/col selection vsparse#50) to stay a genuinely
    lazy view for a duck-typed backend, falling back to ordinary bracket
    indexing for plain dense/sparse X_mat (unaffected, still cheap for
    those). Used at all three X_mat[mask] sites in _bicv_trial.
  • _test_block_moments gains 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 whole
    subset as one dense block -- mirroring the existing scipy-sparse branch's
    own streaming design, just against a different backend's primitives.
  • pyproject.toml: temporarily pins vsparse to its (as yet
    unmerged/unreleased) select-without-recalculate branch -- drop once
    Add 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-existing
    skip, unrelated)
  • uv run ruff check . / uv run ruff format --check .
  • uv run ty check scrise/
  • New tests in test_rank_selection.py: _restrict_rows matches
    bracket 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, streams
    across 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

aarmey and others added 2 commits September 13, 2026 12:24
_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>
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>
@aarmey
aarmey merged commit c9d1daf into main Sep 14, 2026
5 checks passed
@aarmey
aarmey deleted the vsparse-lazy-bicv-blocks branch September 14, 2026 17:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant