Let a container hold a normalized view, and stream it to a device in row blocks

[fishidaho]

Draft. The GPU half needs a decision before review — see What is GPU and what is not below. #42 and #52 were both closed, so the device path may not be wanted at all; the container half stands on its own either way.

Problem

VCSCAnnData.X could hold only VCSCArray/VCSRArray, not a normalized view. Every caller wanting to put one in a container had to go through to_scipy_sparse() first, materializing the whole matrix purely to have somewhere to put it — 42.3 GB on the IBDverse cohort — which defeats the point of a lazy, .select-able view.

Separately, a bulk device copy is not possible above 2**31 nonzeros: CuPy's csr_matrix derives one shared index dtype from the contents, so indices widen to int64 and the cohort needs 42.3 GB (14.1 GB float32 values plus 28.2 GB indices) before any operand is allocated. On a 24.5 GB card that fails at the second allocation.

Fix

X accepts a NormalizedViewBase, and subsetting an AnnData routes through select(..., recalculate=False) — keeping the view's statistics rather than renormalizing the subset, which is what slicing an already-normalized matrix means.

to_device() walks the view in row blocks of roughly chunk_nnz nonzeros instead of uploading the matrix. Every block is individually under 2**31 nonzeros, so its indices stay int32 and device residency is bounded by the block rather than the dataset. Decoded blocks are cached, so repeated passes do not re-decode off the packed view.

Measurements

Cohort BiCV, 2 ranks x 1 repeat:

	wall	host
lazy + CPU	488 s	45.3 GB
chunked GPU, no cache	755 s	45.0 GB
chunked GPU, cached blocks	468 s	47.1 GB

Device demand: 42.3 GB (OOM) → 3.9 GB. The block cache costs ~2 GB of host and buys back the whole regression; it is under 2**31 nonzeros by construction, so its indices stay int32 and it costs about a third less than slicing one int64 matrix would (28.2 GB vs 42.3 GB). Results unchanged throughout.

What is GPU and what is not

The two halves are separable and only one needs a GPU to review.

Container, no GPU — _anndata_class.py, 31 lines. _X_TYPES admits NormalizedViewBase, _subset_2d handles it, _check_vcs_type widened. Reviewable and testable on any machine.

Device path, GPU only — _norm_common.py, 208 lines, almost all of it _DeviceNormalizedView plus to_device(). Every CuPy import is inside a method, so there is no new dependency and nothing imports CuPy unless to_device is called.

If the device path is not wanted, dropping it leaves the container change intact.

Known gaps

• No tests. The container half is testable on CI and should get them before this is un-drafted. The device half cannot be tested without a GPU runner.
• Device peak is 23.8 GB of a 24.5 GB card, and sits above 20 GB for 37% of the run. That is dense intermediates and CuPy's non-returning pool, not the blocks — so this cannot currently share a card. Reusing one device buffer across blocks is the next step.

Merge order

Last, and only after the GPU question is settled.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Rvu3cf8ZL7F5EPX22eo6Je

[fishidaho]

Pushed two fixes found while verifying the branch:

• raw_X no longer accepts a normalized view. Widening X had widened raw_X through the shared type check, which broke six test_raw_x_setter_validation cases. raw_X holds raw counts by definition, so only X takes the wider set now.
• ruff. __slots__ unsorted and two calls wrapped narrower than the line length; the branch had never been linted.

Suite is green on the branch: 1261 passed, ty / ruff / codespell clean. The "no tests" gap in the description still stands — nothing here adds coverage for the new container or device paths.