Allocate the prediction output up front instead of concatenating batches

[RobbinBouwmeester]

Problem

predict() on a multitask model without task_idx returns one column per LC setup. At 6,543 setups that is 25.6 kB per peptide, so a 2.6 M PSM run needs a 63 GiB result.

_predict_epoch built it by collecting every batch in a list and calling torch.cat at the end. That needs the result twice over, once in the parts and once in the copy, so the peak is ~126 GiB. Worse, it only fails after every batch has been predicted.

A user hit this through MS2Rescore, which calls predict(..., return_matrix=True) for all PSMs at once. It failed after 100 minutes of prediction with a bare allocator error:

ERROR [enforce fail at alloc_cpu.cpp:115] data. DefaultCPUAllocator:
not enough memory: you tried to allocate 67731356304 bytes.
  File "deeplc/_model_ops.py", line 365, in _predict_epoch
    return torch.cat(predictions, dim=0)

Nothing in that tells the caller the output was 2,587,932 x 6,543, or that task_idx exists.

Change

Allocate the output once, on the first batch, and fill it in place:

• Halves the peak. One copy instead of the parts plus the concatenation.
• Fails before the compute, not after. The allocation is attempted on the first batch.
• Says what went wrong. The error names the size and points at task_idx.

MemoryError: Predicting would need an output of 2,587,932 x 6,543 values (63.1 GiB).
A multitask model returns one column per LC setup, so ask for the setups you need
with predict_kwargs={'task_idx': [...]}, or predict in smaller groups of peptides.

The bucketed path already preallocated; it now shares the same helper so its failure carries the same message. Datasets that cannot report a length still fall back to collecting batches.

Not addressed here

This makes the failure fast and legible. It does not make the 63 GiB request succeed. FactorHead is rank-64, so that matrix is a rank-64 product plus a per-column affine and could be held in (n, 64) instead, 102x smaller. That is a follow-up PR.

Tests

• preallocated output equals the batchwise concatenation it replaces
• an allocation that cannot be satisfied raises MemoryError naming task_idx
• the hint reports the size in GiB

230 tests pass. ruff check on the touched files reports the same 5 pre-existing findings as main, none added.

🤖 Generated with Claude Code