dataset: split `LeechDataset` into a frozen `TensorCorpus` and an augmenting dataset; block-only twins; batched cross-layer augmentation

[jayhesselberth]

Tier P3 · effort L · from the 2026-09-13 review (training F6, F7; models F9)

LeechDataset (src/leech/dataset.py:510-1941, 1,430 lines; __init__ 369
lines, 21 parameters, 22 fields) is a corpus loader (_ArrayStream,
_TensorFill, row and block fill, CSR expansion, the encoding-coverage
decision), a tensor store, and an augmenter. External code reaches into its
privates: resolve_val_dataloader_workers reads _signals_tensor
(dataset.py:2064), train_model reads train_dataset.chunks[0], benchmark
iterates dataset.chunks. model_type is only used to look up two
hand-maintained name sets. Four rules exist as row/block twins held bit-equal
by tests/test_dataset_streaming.py::TestBlockFillParity:
_prepare_signal/_prepare_signals_block (1415/1267),
_prepare_features/_prepare_features_block (1489/1314),
_encode_sequence/_encode_sequences_block (1665/1253),
_apply_augmentation/_apply_augmentation_batch (1539/1904). Enabling
--augment-shift-max-bases or --augment-time-mask-bases disables
__getitems__ entirely (dataset.py:1856-1857), the measured 9× cliff.
append_dwell_template_channels (dataset.py:186-229) is a Python double
loop with no test and no Rust counterpart.

Acceptance criteria

• TensorCorpus (frozen: one contiguous tensor per field, a metadata ChunkTable, effective_seq_encoding, is_contiguous), built by TensorCorpus.from_npz(...) / .from_chunks(...); LeechDataset(corpus, augment: AugmentConfig | None, *, needs_features, wide_features) only indexes and augments. No external module reads a _private attribute of either (test with a grep-based assertion or __slots__).
• The block form is the implementation; the row form is block(x[None])[0]; TestBlockFillParity becomes golden-fixture tests (saved expected tensors), so the reference is data, not a second implementation. _block_fill_supported shrinks accordingly.
• Cross-layer shift and time mask draw (B,) parameters and apply with gather/mask tensors on the batched path; __getitems__ is never disabled by an augmentation.
• append_dwell_template_channels is vectorised and tested, and tests/test_backend_parity.py notes it as Python-only (or it gains a Rust twin).
• The fork-shares-contiguous-buffers invariant and the numpy-gather rule are preserved and asserted by the existing tests; peak RSS at load on the fixture is unchanged or lower (record it).

Non-goals

• Changing any tensor value: test_train_epoch_matches_host_side_reference and the streaming parity tests must pass unchanged.