Dataset Module¶

def __init__(
    self,
    chunk_path: Path | None = None,
    signal_len: int = 400,
    kmer_len: int = 11,
    model_type: str = "ConvLSTMDwell",
    dwell_offset: int = 0,
    chunks: list[dict] | None = None,
    augmentation: dict | None = None,
    seq_encoding: str = "signal_kmer",
    signal_kmer_context: tuple[int, int] = (4, 4),
    allow_encoding_fallback: bool = True,
    left_context: int | None = None,
    right_context: int | None = None,
    strict_window: bool = False,
    confound_encoder: "ConfoundEncoder | None" = None,
    label_noise_rates: dict[str, float] | None = None,
    cl_regression: bool = False,
    signal_mode: str = "both",
    time_mask_bases: int = 0,
    time_mask_count: int = 1,
    shift_max_bases: float = 0.0,
    feature_noise_scale: float = 0.0,
    time_stretch_range: tuple[float, float] | None = None,
    dwell_template_table: str | Path | None = None,
    motif_offset: int | None = None,
    standardize_features: bool = False,
    feature_channel_names: Sequence[str] | None = None,
    selected_feature_channels: Sequence[str] | None = None,
):
    """
    Initialize dataset.

    Args:
        chunk_path: Path to .npz file with training chunks
        signal_len: Expected signal length (will pad/truncate)
        kmer_len: Expected k-mer length
        model_type: Model architecture name (e.g., "ConvLSTMDwell", "TransformerDwell")
        dwell_offset: Shift dwell/feature window toward 3' end (bases).
            Compensates for physical offset between motor protein and
            sensing region. Requires feature_left >= kmer_context + offset.
        chunks: Pre-loaded list of chunk dicts. When provided, chunk_path is
            ignored and no disk I/O occurs (useful for grid search caching).
        augmentation: Signal augmentation config dict. Keys:
            - jitter_std (float): Gaussian noise std dev (0 = disabled)
            - scale_range (tuple[float, float]): Random scale range (1.0, 1.0 = disabled)
        allow_encoding_fallback: When ``seq_encoding="signal_kmer"`` and the
            corpus carries no base-to-signal maps at all, fall back to
            ``base_onehot`` (warning) rather than raising. A corpus that
            carries them for *some* chunks is damaged and raises either way.
            Read :attr:`effective_seq_encoding` for what the dataset
            actually yields — it is a different model input, not a tuning
            difference, so callers that persist a config must record it.
        left_context: Left signal context (samples before focus base).
            When both left_context and right_context are provided, crop
            asymmetrically around the focus base instead of center-cropping.
        right_context: Right signal context (samples after focus base).
        strict_window: When the asymmetric crop (``left_context``/
            ``right_context``) reaches outside the stored chunk -- the
            corpus was prepared with a narrower ``signal_context`` than
            this window -- the shortfall is zero-padded either way; the
            default (False) logs one warning for the whole dataset the
            first time it happens, ``True`` raises immediately instead.
            A silently zero-padded tail is exactly what went undetected
            in two production retrains that trained with a wider
            ``right_context`` than the corpus was prepared with (#255).
        confound_encoder: Optional :class:`~leech.confounds.ConfoundEncoder`
            that maps each chunk to an integer confound class. When provided,
            each batch includes a ``confound_label`` tensor for adversarial
            training. Chunks whose confound value is unknown get ``-1``
            (ignored by the CE loss via ``ignore_index=-1``).
        label_noise_rates: Optional ``{source_group: flip_rate}`` mapping.
            When provided, each batch includes a ``noise_rate`` tensor
            (shape matching ``label``) built by looking up each chunk's
            ``source_group`` (missing/empty -> "unknown"); groups absent
            from the mapping get rate 0. Consumed by
            ``leech.losses.NoiseCorrectedBCEWithLogitsLoss``.
        cl_regression: When True, include ``cl_target`` in each batch
            (cl_value / 255.0 in [0,1]; sentinel -1.0 for missing).
        time_mask_bases: Max width in bases for time masking (0 = disabled).
            Zeros out contiguous blocks across signal, features, and sequence.
        time_mask_count: Number of time masks to apply per sample.
        shift_max_bases: Max cross-layer shift in bases (0 = disabled).
            Simulates motif anchor offset, applied consistently to all branches.
        feature_noise_scale: Per-channel Gaussian noise multiplier (0 = disabled).
            Noise std = feature_noise_scale * per-channel empirical std.
        time_stretch_range: ``(min, max)`` per-sample time-stretch factor,
            drawn uniformly per chunk ``(1.0, 1.0)`` = disabled). Resamples
            the signal window about the focus position (cropped/padded back
            to ``signal_len``), scales the CSR base-to-signal map by the
            same factor (so ``signal_kmer`` stays aligned), and scales the
            dwell-derived feature channels -- see
            :meth:`_apply_time_stretch` for the exact per-channel rule.
            Level channels are unchanged. Implemented only on the batched
            ``__getitems__`` fetch path: incompatible with
            ``shift_max_bases``/``time_mask_bases`` (which force the
            per-sample fallback) and with a corpus that degrades to a
            per-chunk list -- both raise at construction rather than
            silently skipping the augmentation.
        dwell_template_table: Path to TSV with per-AA per-position expected
            dwell times. When provided, 20 dwell ratio channels are appended
            to features: ``dwell / expected_dwell[AA_i, pos]`` for each of
            20 AAs. The correct AA's channel has ratio closest to 1.0.
            TSV columns: aa, position, dwell_mean, ... -- see
            :func:`load_dwell_template_table` for the optional header lines
            recording the motif_offset the table was built against. The
            ``dwell`` row to divide is resolved by name from
            :attr:`effective_feature_channels`, not assumed to be row 0
            (issue #366) -- raises ``ValueError`` at construction if
            ``dwell`` is not an effective channel. The table's covered
            position range is also checked against the corpus's feature
            window (issue #367; see :func:`check_dwell_template_coverage`
            and :attr:`dwell_template_positions`/
            :attr:`dwell_template_uncovered_columns`). Incompatible with
            ``time_stretch_range`` (issue #368): the template ratios are
            computed once here and are not rescaled by time-stretch.
        motif_offset: This run's motif_offset, used only to guard
            ``dwell_template_table`` against issue #364 -- a table whose
            ``position`` column was counted from a different focus base.
            Checked via :func:`check_dwell_template_motif_offset` when
            both this and ``dwell_template_table`` are given; ``None``
            (default) skips the check, e.g. for callers that do not yet
            know their own motif_offset.
        standardize_features: Compute corpus-wide per-channel feature
            mean/std once (see :attr:`feature_mean`/:attr:`feature_std`)
            for ``--standardize-features``. This dataset never applies
            them to its own output -- the caller (``train_model``) reads
            them and passes them to the model constructor, which builds a
            frozen affine layer from them. Meaningful only for the
            *training* dataset; a validation/test/predict dataset should
            leave this False and reuse the training corpus's stats via
            the model instead.
        feature_channel_names: The corpus's recorded feature-channel list
            (issue #358), e.g. read from ``prepare_config.json``'s
            ``feature_channels`` key. ``None`` (default) means
            :data:`~leech.constants.DEFAULT_FEATURE_CHANNELS` -- the
            historical twelve, which keys ``_apply_time_stretch``'s
            dwell/dwell_log/dwell_mean rule to the same positions (0, 1,
            2) it always used, so an unset corpus behaves identically to
            before this existed. Consulted when ``time_stretch_range`` is
            active (for the by-name lookup) and whenever
            ``selected_feature_channels`` is given (as the *nominal* list
            to reconcile against the tensor's actual stored row count --
            see :func:`~leech.constants.resolve_stored_feature_channels`).
        selected_feature_channels: Channel ablation (issue #358
            follow-up) -- select and reorder a SUBSET of the corpus's
            feature rows by name, resolved once here against the actual
            stored row count (not assumed position) and applied to every
            per-chunk feature array before anything else touches it, so
            ``num_features``, per-channel standardization stats,
            feature-noise augmentation and time-stretch's by-name dwell
            lookup all see the selected tensor. ``None`` (default) keeps
            every row exactly as before this existed -- selection logic
            does not run at all in that case, so an unset corpus is
            byte-identical. Raises ``ValueError`` naming any requested
            channel absent from the corpus's actual data (see
            :func:`~leech.constants.select_feature_channels`).
    """
    self.chunk_path = chunk_path
    self.signal_len = signal_len
    self.kmer_len = kmer_len
    self.model_type = model_type
    self.dwell_offset = dwell_offset
    self.augmentation = augmentation
    self.seq_encoding = seq_encoding
    self.signal_kmer_context = signal_kmer_context
    self.left_context = left_context
    self.right_context = right_context
    self._strict_window = strict_window
    self._crop_pad_warned = False
    self._time_mask_bases = time_mask_bases
    self._time_mask_count = time_mask_count
    self._shift_max_bases = shift_max_bases
    self._feature_noise_scale = feature_noise_scale
    self._standardize_features = standardize_features
    # Name -> row index in the recorded feature-channel list (issue
    # #358), resolved once here rather than per-batch. `None` when a
    # name is absent from the list (a custom `--feature-channels` run
    # that dropped it, or a channel that needs a k-mer table the run
    # didn't configure) -- `_apply_time_stretch` skips the corresponding
    # transform rather than guessing at a position.
    _fc_names = (
        list(feature_channel_names)
        if feature_channel_names is not None
        else list(DEFAULT_FEATURE_CHANNELS)
    )
    self._dwell_idx = _fc_names.index("dwell") if "dwell" in _fc_names else None
    self._dwell_log_idx = _fc_names.index("dwell_log") if "dwell_log" in _fc_names else None
    self._dwell_mean_idx = _fc_names.index("dwell_mean") if "dwell_mean" in _fc_names else None
    self._time_stretch_range = (
        tuple(time_stretch_range) if time_stretch_range is not None else (1.0, 1.0)
    )
    if self._time_stretch_range != (1.0, 1.0):
        lo, hi = self._time_stretch_range
        if lo <= 0 or hi <= 0 or lo > hi:
            raise ValueError(
                f"time_stretch_range must be a positive (min, max) with min <= max, "
                f"got {self._time_stretch_range}"
            )
        if self._forces_per_sample_fetch():
            raise ValueError(
                "time_stretch_range is only implemented on the batched __getitems__ "
                "fetch path, which shift_max_bases and time_mask_bases both bypass "
                "(they draw one offset per sample and fall back to the per-sample "
                "__getitem__ path). Disable one or the other."
            )
        if dwell_template_table is not None:
            # The 20 template ratio rows are computed once, at
            # construction, as dwell / expected -- time-stretch then
            # rescales `dwell` per sample without touching them, so under
            # stretch they silently stop equalling dwell / expected
            # (issue #368). Refuse the combination rather than accept it
            # quietly, same shape as the shift_max_bases/time_mask_bases
            # refusal just above.
            raise ValueError(
                "time_stretch_range is incompatible with dwell_template_table: "
                "the precomputed dwell-template ratio channels are not rescaled "
                "by time-stretch, so they would silently disagree with a "
                "stretched 'dwell' channel. Disable one or the other."
            )

    # Load dwell template table for 20-channel AA template features
    self._dwell_templates: np.ndarray | None = None
    self._template_min_pos: int = 0
    self._template_max_pos: int = 0
    self._template_aa_order: list[str] = []
    self._dwell_template_table_path: Path | None = None
    self.motif_offset = motif_offset
    self.dwell_template_motif_offset: int | None = None
    self.dwell_template_motif: str | None = None
    self.dwell_template_positions: tuple[int, int] | None = None
    self.dwell_template_header: dict[str, str] = {}
    self.dwell_template_uncovered_columns: int | None = None
    if dwell_template_table is not None:
        self._load_dwell_templates(Path(dwell_template_table))

    self._needs_features = requires_features(model_type)
    # Cached like _needs_features above: for a TOML/Graph architecture,
    # wide_features() re-resolves the config's `[params]` (dict build +
    # eval of every "${...}" expression) on every call, and the
    # row-wise fallback path calls it once per chunk.
    self._wide_features = wide_features(model_type)

    # Use pre-loaded chunks or load from file. Loading from a path streams
    # the per-chunk arrays out of the npz a row block at a time instead of
    # holding a full numpy copy alongside the tensors built from it (#211);
    # pre-loaded chunks have already paid that cost.
    self._array_stream: _ArrayStream | None = None
    self._npz_members: set[str] = set()
    self._s2s_csr: tuple[np.ndarray, np.ndarray] | None = None
    self._s2s_rows: np.ndarray | None = None
    if chunks is not None:
        logger.info(f"Using {len(chunks)} pre-loaded chunks (skipping disk I/O)")
        self.chunks = chunks
    elif chunk_path is not None:
        self._load_from_path(
            Path(chunk_path), signal_mode=signal_mode, seq_encoding=seq_encoding
        )
    else:
        raise ValueError("Either chunk_path or chunks must be provided")

    # Filter chunks with valid numeric labels (label_int). The columnar
    # path applied this at load, recording the same mask so npz rows stay
    # aligned with self.chunks — a mismatch here would pair signals with
    # the wrong labels silently.
    if not isinstance(self.chunks, ChunkTable):
        self.chunks = [c for c in self.chunks if c["label_int"] is not None]

    if len(self.chunks) == 0:
        raise ValueError(f"No valid chunks found{f' in {chunk_path}' if chunk_path else ''}")

    # Channel ablation (issue #358 follow-up): resolve and apply a
    # `selected_feature_channels` subset, once, against the corpus's
    # ACTUAL stored row count -- not the (possibly wider, if a `kmer_*`
    # name was recorded but never computed) nominal recorded list. Gated
    # entirely on `selected_feature_channels is not None` so an unset
    # corpus runs none of this new code and stays byte-identical to
    # before it existed.
    self._feature_row_idx: np.ndarray | None = None
    self.effective_feature_channels: list[str] | None = None
    if self._needs_features:
        if selected_feature_channels is not None:
            actual_count = self._stored_feature_row_count()
            if actual_count is None:
                raise ValueError(
                    "selected_feature_channels was requested but this dataset's "
                    "stored feature row count could not be determined"
                    f"{f' ({chunk_path})' if chunk_path else ''}"
                )
            available = resolve_stored_feature_channels(_fc_names, actual_count)
            row_idx = select_feature_channels(available, selected_feature_channels)
            self._feature_row_idx = np.asarray(row_idx, dtype=np.int64)
            effective_names = list(selected_feature_channels)
            # The tensor's channel axis now has the SELECTED order, not
            # the corpus's recorded order -- recompute the by-name
            # indices _apply_time_stretch uses against it.
            self._dwell_idx = (
                effective_names.index("dwell") if "dwell" in effective_names else None
            )
            self._dwell_log_idx = (
                effective_names.index("dwell_log") if "dwell_log" in effective_names else None
            )
            self._dwell_mean_idx = (
                effective_names.index("dwell_mean") if "dwell_mean" in effective_names else None
            )
            self.effective_feature_channels = effective_names
        else:
            self.effective_feature_channels = _fc_names

    # Resolve the dwell-template append's row index by NAME, not by
    # convention (issue #366): a --feature-channels reorder can put
    # `dwell` at any row, and dividing row 0 regardless silently computes
    # nonsense ratios for every non-default channel order. Also check the
    # table's position coverage of the corpus's feature window, once,
    # here rather than per chunk (issue #367).
    self._dwell_template_row: int | None = None
    if self._dwell_templates is not None:
        if self.effective_feature_channels is None or (
            "dwell" not in self.effective_feature_channels
        ):
            raise ValueError(
                "dwell_template_table requires a 'dwell' feature channel, but "
                f"the effective feature channels are "
                f"{self.effective_feature_channels!r}"
            )
        self._dwell_template_row = self.effective_feature_channels.index("dwell")
        logger.info(
            "dwell-template append resolved dwell channel to row %d of %s",
            self._dwell_template_row,
            self.effective_feature_channels,
        )

        kmer_context = self.kmer_len // 2
        fw_source = self.chunks if isinstance(self.chunks, ChunkTable) else self.chunks[0]
        raw_feat_start, raw_feat_end = feature_window_from_metadata(fw_source, kmer_context)
        feat_start, feat_end, _ = resolve_feature_window(
            raw_feat_start, raw_feat_end, kmer_context
        )
        self.dwell_template_uncovered_columns = check_dwell_template_coverage(
            self._dwell_template_table_path,
            self._template_min_pos,
            self._template_max_pos,
            feat_start,
            feat_end,
        )

    # Pre-tensorize: encode sequences, labels, signals, and features once.
    # Each accumulator fills one preallocated contiguous tensor — see
    # _TensorFill for why stacking a list instead doubles the peak.
    n_chunks = len(self.chunks)
    fill_encoded_seqs = _TensorFill("Encoded-sequence", n_chunks)
    fill_labels = _TensorFill("Label", n_chunks)
    fill_signals = _TensorFill("Signal", n_chunks)
    fill_features = _TensorFill("Feature", n_chunks)
    fill_confounds = _TensorFill("Confound", n_chunks)
    fill_cl_targets = _TensorFill("CL target", n_chunks)
    fill_noise_rates = _TensorFill("Noise rate", n_chunks)
    fill_focus_idx = _TensorFill("Focus-idx", n_chunks)
    self._encoded_seqs: list[torch.Tensor] = []
    self._labels: list[torch.Tensor] = []
    self._signals: list[torch.Tensor] = []
    self._features: list[torch.Tensor] = []
    self._confound_encoder = confound_encoder
    self._has_confound = confound_encoder is not None
    self._confound_labels: list[torch.Tensor] = []
    self._cl_regression = cl_regression
    self._cl_targets: list[torch.Tensor] = []
    self._label_noise_rates = label_noise_rates or {}
    self._has_noise_rates = bool(label_noise_rates)
    self._noise_rates: list[torch.Tensor] = []

    # Determine effective encoding, from the whole corpus rather than from
    # chunk 0 (#230). A corpus carries the signal_kmer inputs for every
    # chunk or for none: none is the version-skew case the fallback exists
    # for, and anything in between is damage.
    self._effective_seq_encoding = seq_encoding
    if seq_encoding == "signal_kmer":
        n_covered, n_total = self._signal_kmer_coverage()
        if n_covered < n_total:
            n_missing = n_total - n_covered
            detail = (
                f"{n_missing} of {n_total} chunks ({n_missing / n_total:.1%}) lack "
                f"seq_to_sig_map/sequence_with_kmer_context"
                f"{f' in {chunk_path}' if chunk_path else ''}"
            )
            if n_covered:
                # Ragged, not legacy — and no automatic choice is right.
                # Encoding the uncovered rows anyway gives them all-zero
                # sequence channels; switching the whole corpus to
                # base_onehot on the strength of a few bad rows throws the
                # encoding away for every good one. Both are the silent
                # representation change this guard exists to prevent.
                raise ValueError(
                    f"seq_encoding='signal_kmer': {detail}. A corpus carries "
                    "base-to-signal maps for every chunk or for none; re-prepare "
                    "it, or ask for --seq-encoding base_onehot explicitly."
                )
            if not allow_encoding_fallback:
                raise ValueError(
                    f"seq_encoding='signal_kmer' was requested but {detail}. "
                    "Re-prepare the corpus with a leech version that writes "
                    "base-to-signal maps, or ask for --seq-encoding base_onehot "
                    "explicitly."
                )
            logger.warning("%s; falling back to base_onehot encoding", detail)
            self._effective_seq_encoding = "base_onehot"
            self._s2s_csr = None

    # Detect signal_residual channel and apply signal_mode
    self._signal_mode = signal_mode
    if self._array_stream is not None:
        # Arrays were deferred, so presence is a property of the file.
        self._has_signal_residual = bool(
            {"signal_residuals_flat", "signal_residuals"} & self._npz_members
        )
    else:
        self._has_signal_residual = self.chunks[0].get("signal_residual") is not None
    if signal_mode == "both" and self._has_signal_residual:
        self.signal_channels = 2
    else:
        self.signal_channels = 1

    # Detect multi-class: if any label_int > 1, use long dtype for
    # CrossEntropyLoss. Off the column when there is one — the generator
    # form builds a row view per chunk, 150 ms per 200k chunks for a max.
    if isinstance(self.chunks, ChunkTable):
        max_label = int(self.chunks.require_values("label_int").max())
    else:
        max_label = max(c["label_int"] for c in self.chunks)
    self._multiclass = max_label > 1

    # For signal_kmer encoding, the on-the-fly inputs are a few hundred
    # bytes/chunk (seq_ints int8 + seq_to_sig_map int32) vs ~77 KB/chunk
    # for the encoded output — orders of magnitude smaller. We stash the
    # compact inputs here and encode lazily: `__getitems__` calls
    # `encode_signal_kmer_batch` once per DataLoader batch (rayon-parallel
    # in Rust, issue #260) rather than looping the single-row encoder once
    # per chunk; `__getitem__`'s single-item path still calls the
    # single-row `encode_signal_kmer`. base_onehot stays eagerly
    # tensorized (88 floats/chunk; not worth deferring).
    self._seq_ints: list[np.ndarray] = []
    self._seq_to_sig: list[np.ndarray] = []
    self._seq_ints_tensor: torch.Tensor | None = None

    # Arrays come either from the chunk dicts (pre-loaded / legacy corpus)
    # or a row-block stream over the npz. Both fill the same accumulators.
    # A streamed corpus is handled a whole block at a time: the arrays
    # arrive in blocks of ~1,500 rows, and taking them apart to prepare one
    # row at a time cost 58-67 us per chunk — six to seven minutes of
    # single-threaded startup on a 6.7M-chunk corpus, none of it I/O.
    fills = {
        "signals": fill_signals,
        "encoded_seqs": fill_encoded_seqs,
        "labels": fill_labels,
        "features": fill_features,
        "confounds": fill_confounds,
        "cl_targets": fill_cl_targets,
        "noise_rates": fill_noise_rates,
    }
    # Per-chunk focus index within the final signal_len tensor, needed only
    # by time-stretch. Computed here (not re-derived later) because the raw,
    # pre-pad/crop signal length is only available during this fill pass --
    # self.chunks' array fields are nulled out once tensorization finishes.
    if self._time_stretch_range != (1.0, 1.0):
        fills["focus_idx"] = fill_focus_idx
    if self._block_fill_supported():
        self._fill_from_blocks(fills)
    else:
        self._fill_from_rows(fills)

    # Every per-chunk tensor now lives in one contiguous buffer. That isn't
    # just for cache friendliness — it's required for fork-safety. A
    # DataLoader with num_workers > 0 forks worker processes that COW-inherit
    # the parent's address space. CPython refcounts live inside each
    # PyObject header, so a worker iterating a list of N tensors writes to N
    # separate page-resident headers and faults every page into a private
    # copy, multiplying peak RSS by (1 + num_workers). A single contiguous
    # tensor keeps its data buffer outside Python's GC, so the buffer
    # pages are genuinely shared across the fork.
    self._signals_tensor, self._signals = fill_signals.finish()
    self._encoded_seqs_tensor, self._encoded_seqs = fill_encoded_seqs.finish()

    # Compact signal_kmer inputs — only populated when encoding == signal_kmer.
    # Stacked into fork-safe int tensors. encode_signal_kmer is then called
    # lazily per-sample in __getitem__ (workers parallelize behind prefetch).
    # Chunks have variable basecalled sequence lengths near the motif, so we
    # pad to the per-array max. Padding values are chosen so the encoder
    # gracefully ignores them: seq_ints=-1 hits the encoder's `base < 0`
    # skip; seq_to_sig=signal_len makes its `sig_start < signal_len` check
    # fail, writing nothing for the padded positions.
    #
    # The block-wise filler builds the padded int8 matrix straight off the
    # text column and leaves ``_seq_ints_tensor`` set; the per-chunk path
    # leaves a list to pad here. Either way the CSR expansion below runs.
    self._seq_to_sig_tensor: torch.Tensor | None = None
    if self._effective_seq_encoding == "signal_kmer" and (
        self._seq_ints or self._seq_ints_tensor is not None
    ):
        if self._seq_ints:
            max_seq_ints_len = max(s.shape[0] for s in self._seq_ints)
            n = len(self._seq_ints)
            padded_seq_ints = np.full((n, max_seq_ints_len), -1, dtype=np.int8)
            for i, si in enumerate(self._seq_ints):
                padded_seq_ints[i, : si.shape[0]] = si
            self._seq_ints_tensor = torch.from_numpy(padded_seq_ints)
            self._seq_ints = []
        n = int(self._seq_ints_tensor.shape[0])

        if self._s2s_csr is not None:
            # Streaming path: expand every map at once from the CSR pair,
            # instead of one astype/clip per chunk.
            values, offsets = self._s2s_csr
            crop_starts = None
            if self.left_context is not None and self.right_context is not None:
                crop_starts = self._crop_starts(values, offsets) - self.left_context
            padded_s2s = _expand_seq_to_sig_csr(
                values,
                offsets,
                self._s2s_rows,
                signal_len=signal_len,
                crop_starts=crop_starts,
            )
            self._s2s_csr = None
        else:
            max_s2s_len = max(s.shape[0] for s in self._seq_to_sig)
            padded_s2s = np.full((n, max_s2s_len), signal_len, dtype=np.int32)
            for i, s2s in enumerate(self._seq_to_sig):
                padded_s2s[i, : s2s.shape[0]] = s2s
        self._seq_to_sig_tensor = torch.from_numpy(padded_s2s)
        self._seq_to_sig = []

    self._labels_tensor, self._labels = fill_labels.finish()

    self._features_tensor: torch.Tensor | None = None
    if self._needs_features:
        self._features_tensor, self._features = fill_features.finish()

    self._confound_labels_tensor: torch.Tensor | None = None
    if self._has_confound:
        self._confound_labels_tensor, self._confound_labels = fill_confounds.finish()

    self._cl_targets_tensor: torch.Tensor | None = None
    if self._cl_regression:
        self._cl_targets_tensor, self._cl_targets = fill_cl_targets.finish()

    self._noise_rates_tensor: torch.Tensor | None = None
    if self._has_noise_rates:
        self._noise_rates_tensor, self._noise_rates = fill_noise_rates.finish()

    self._focus_idx_tensor: torch.Tensor | None = None
    if self._time_stretch_range != (1.0, 1.0):
        self._focus_idx_tensor, _leftover_focus_idx = fill_focus_idx.finish()
        if _leftover_focus_idx:
            # A per-chunk scalar column never disagrees in shape, so this
            # should be unreachable -- but if it ever is, time-stretch has
            # no per-row focus to anchor on, and _apply_time_stretch would
            # silently fall back to a wrong constant. Fail loudly instead.
            raise ValueError(
                "time_stretch_range: could not build a per-chunk focus-index "
                "column (unexpected shape mismatch in _TensorFill)"
            )

    # Drop the raw numpy arrays from self.chunks now that everything has
    # been pre-tensorized. External code (samplers, label tally, feature
    # window introspection) still reads the small scalar/string fields,
    # so we keep self.chunks alive but null out the per-chunk arrays.
    # Without this, each chunk dict keeps a ~50 KB numpy view alive and
    # the same COW blowup hits during DataLoader fork.
    if not isinstance(self.chunks, ChunkTable):
        for chunk in self.chunks:
            for key in (
                "signal",
                "signal_residual",
                "dwell",
                "features",
                "seq_to_sig_map",
                "sequence_with_kmer_context",
            ):
                if key in chunk:
                    chunk[key] = None

    # Same reasoning for the streaming bookkeeping: one row per chunk each,
    # and a DataLoader that spawns workers pickles whatever is still here.
    self._s2s_csr = None
    self._s2s_rows = None
    if self._array_stream is not None:
        self._array_stream.keep = None

    # Precompute per-channel feature stds for feature noise augmentation.
    # Reuse the already-stacked features tensor when available.
    self._feature_stds: torch.Tensor | None = None
    if self._feature_noise_scale > 0 and self._needs_features:
        if self._features_tensor is not None:
            self._feature_stds = self._features_tensor.std(dim=0)  # (C, K)
        else:
            logger.warning("Feature shapes differ, feature noise disabled")
            self._feature_noise_scale = 0.0

    # Corpus-wide per-channel mean/std for --standardize-features (#283):
    # one number per feature channel, reusing the same already-stacked
    # _features_tensor as the noise pass above rather than a second scan.
    # Collapsed over both the chunk axis AND the window axis (dim 0 and
    # 2) -- unlike the noise std, which is deliberately per-position --
    # because a channel's physical scale (a dwell sample count, a
    # MAD-unit level stat, ...) is a property of the channel, not of
    # where in the window it's read.
    self.feature_mean: torch.Tensor | None = None
    self.feature_std: torch.Tensor | None = None
    if self._standardize_features and self._needs_features:
        if self._features_tensor is not None:
            self.feature_mean = self._features_tensor.mean(dim=(0, 2))  # (C,)
            self.feature_std = self._features_tensor.std(dim=(0, 2))  # (C,)
        else:
            logger.warning(
                "Feature shapes differ across chunks; --standardize-features disabled"
            )
            self._standardize_features = False

    # Approx samples per base for cross-layer shift/mask
    self._samples_per_base = signal_len / max(kmer_len, 1)

    # Whether __getitems__ can gather a batch straight out of the
    # contiguous tensors. Every field has to be one: a field that degraded
    # to a list of per-chunk tensors has no batch to gather.
    self._batched_fetch = (
        self._signals_tensor is not None
        and self._labels_tensor is not None
        and (not self._needs_features or self._features_tensor is not None)
        and (not self._has_confound or self._confound_labels_tensor is not None)
        and (not self._cl_regression or self._cl_targets_tensor is not None)
        and (not self._has_noise_rates or self._noise_rates_tensor is not None)
        and (self._time_stretch_range == (1.0, 1.0) or self._focus_idx_tensor is not None)
        and (
            (self._seq_ints_tensor is not None and self._seq_to_sig_tensor is not None)
            if self._effective_seq_encoding == "signal_kmer"
            else self._encoded_seqs_tensor is not None
        )
    )
    if self._time_stretch_range != (1.0, 1.0) and not self._batched_fetch:
        raise ValueError(
            "time_stretch_range requires the dataset to support the batched "
            "__getitems__ fetch path (contiguous per-field tensors); this corpus "
            "degraded to a per-chunk list (see the 'shapes differ' warning above "
            "if logged), which has no per-sample implementation of time-stretch."
        )

    _n_encoded = (
        self._encoded_seqs_tensor.shape[0]
        if self._encoded_seqs_tensor is not None
        else len(self._encoded_seqs)
    )
    logger.debug(
        f"Pre-tensorized {len(self.chunks)} chunks "
        f"({_n_encoded} sequences encoded, encoding={self._effective_seq_encoding})"
    )