giant/giant/sample.py

import torch
import torch.nn.functional as F

from giant.constants import CONT_SLOT_DIM, X_DIM
from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim
from giant.model.schedule import CosineSchedule


def _predict_n_sec_if_owned(
    model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor | None:
    """Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
    (fresh runs move it to stage 2 — see `Stage1Model`'s docstring). `None`
    here means "ask stage 2 instead", which every caller (`giant/rollout.py`,
    `giant/cli.py`) must do for a fresh checkpoint."""
    if getattr(model, "n_sec_head", None) is None:
        return None
    logits = model.predict_n_sec(cond_cont, cond_cat)
    return logits.argmax(dim=-1)


@torch.no_grad()
def sample_flow(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Euler integration of the Stage-1 vector field from t=0 to t=1.

    Returns (primary_sample, n_sec_pred):
        primary_sample: (B, X_DIM) — normalised 9D primary post-step output
        n_sec_pred: (B,) int64 — predicted secondary count, or `None` if
            `model` has no `n_sec_head` (a fresh v0.3.0 Stage1Model — see
            `_predict_n_sec_if_owned`).
    """
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((B,), i * dt, device=device)
        v = model(x, cond_cont, cond_cat, t=t)
        x = x + v * dt
    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)


@torch.no_grad()
def sample_ddpm(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)
    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    T = schedule.T
    for i in reversed(range(T)):
        t_norm = torch.full((B,), i / T, device=device)
        eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
        beta = schedule.betas[i]
        alpha = schedule.alphas[i]
        alpha_bar = schedule.alpha_bars[i]
        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
        x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)


@torch.no_grad()
def sample_ddim(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
    steps: int = 50,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)
    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    T = schedule.T
    timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
    x = torch.randn(B, X_DIM, device=device)
    for step_idx, ts in enumerate(timesteps):
        t_idx = int(ts.item())
        t_norm = torch.full((B,), t_idx / T, device=device)
        eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
        ab_t = schedule.alpha_bars[t_idx]
        if step_idx + 1 < len(timesteps):
            ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
        else:
            ab_prev = torch.ones(1, device=device)
        x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
        x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)


@torch.no_grad()
def sample_wgan(
    generator: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)
    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
    generator.eval()
    B = cond_cont.size(0)
    z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
    x = generator(z, cond_cont, cond_cat)
    return x, _predict_n_sec_if_owned(generator, cond_cont, cond_cat)


def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
    """The width of `sec_decoder`'s own trunk in/out vector — folded
    (continuous + type) under `particle_type.target = "physical"` or
    `generator = "wgan"`, continuous-only otherwise (the type slice then
    comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s
    docstring)."""
    return stage2_trunk_sec_dim(
        sec_decoder.particle_type_cfg,
        sec_decoder.generator_kind,
        sec_decoder.k_max,
        sec_decoder.type_dim,
    )


def _type_folded(sec_decoder: torch.nn.Module) -> bool:
    target = sec_decoder.particle_type_cfg.target
    return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice


def _decode_stage2_flat(
    sec_decoder: torch.nn.Module,
    x: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
    tensors, generator/`particle_type.target`-agnostic: shared by
    `sample_secondaries`/`sample_secondaries_wgan`, which differ only in how
    `x` was produced.

    Returns (sec_cont, sec_type, sec_valid):
        sec_cont:  (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir]
        sec_type:  (B, k_max, type_dim) — under `target="physical"` this is
            [log_mass, charge] (normalised iff the checkpoint's sec_phys
            normalizer was applied at training time — denormalize before
            treating as physical units; see
            giant.data.transforms.decode_secondaries); under `"onehot"` /
            `"embedding"` it is raw class logits / an embedding-space vector
            — decode via giant.particles.decode_topn_class /
            decode_embedding_nearest (see giant/rollout.py).
        sec_valid: (B, k_max) bool — True for slots i < n_sec_pred
    """
    B = x.size(0)
    device = x.device
    k_max = sec_decoder.k_max
    type_dim = sec_decoder.type_dim
    if _type_folded(sec_decoder):
        x_slots = x.view(B, k_max, CONT_SLOT_DIM + type_dim)
        sec_cont = x_slots[:, :, :CONT_SLOT_DIM]
        sec_type = x_slots[:, :, CONT_SLOT_DIM:]
    else:
        sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
        sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
    return sec_cont, sec_type, sec_valid


@torch.no_grad()
def sample_secondaries(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor,
    steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Euler integration of `Stage2OneShot`'s (flow/ddpm) vector field; return
    raw slot outputs — see `_decode_stage2_flat`'s docstring for the returned
    (sec_cont, sec_type, sec_valid) shapes/meaning.

    n_sec_pred: (B,) int64 — number of valid secondaries per step
    """
    sec_decoder.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    flat_width = _stage2_flat_width(sec_decoder)

    x = torch.randn(B, flat_width, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((B,), i * dt, device=device)
        v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
        x = x + v * dt

    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)


@torch.no_grad()
def sample_secondaries_wgan(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Single-pass `Stage2OneShot` WGAN generator sample; see
    `_decode_stage2_flat`'s docstring for the returned (sec_cont, sec_type,
    sec_valid) shapes/meaning."""
    sec_decoder.eval()
    B = cond_cont.size(0)
    z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
    x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)


@torch.no_grad()
def sample_secondaries_ar(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor | None,
    steps: int = 10,
    full_length: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """`Stage2Autoregressive` inference loop: one token at a time, in
    descending-energy slot order, up to `k_max` sequential calls. Unlike
    training (teacher forcing — a single parallel pass over ground-truth
    tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`),
    there is no ground truth at inference: each token's conditioning is built
    free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
    train/inference gap that is the cost of markov history's
    expressiveness.

    A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
    the "K sequential forwards" cost applies per-token here, not once, so a
    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
    time), not `k_max * steps` — see the row-compaction paragraph below.

    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
    `("head", "truth")`, or a stop-token decoder driven by
    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
    run the *full* `k_max`-length free-running self-sample regardless of the
    decoder's own stop head — the scheduled-sampling training contract does
    not truncate, see `full_length` below) fixes each row's secondary count
    up front.

    `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
    (`n_sec.mode = "stop_token"`): before generating each slot's token, that
    slot's own stop logit (`predict_stop`, evaluated on the same prefix
    conditioning as the token itself — see `predict_type`'s docstring for
    why this needs no extra state) decides whether generation should have
    already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
    0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
    `n_sec_pred` is the first slot index where this fires. A row that never
    fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
    modeling ceiling).

    **Row compaction.** A row that has already produced its `n_sec_pred`
    tokens (or, under `stop_token`, has already fired its stop logit) has
    nothing left to contribute — every later slot of that row is masked out
    of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
    `giant/cli.py`) never read it. So unless `full_length=True`, this
    function drops such rows from the active set entirely instead of running
    the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
    (or, under `stop_token`, every row — the first stop decision can fire at
    slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
    `O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
    zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
    what a full, uncompacted run would have written there before masking.
    `AttentionHistory`'s KV cache is kept aligned to the shrinking active set
    via `Stage2Autoregressive.select_history_cache`
    (`giant.model.history.HistoryEncoder.select_cache`) every time the set
    shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
    directly rather than through a cache) is compacted the same way.

    `full_length=True` disables all of the above: every row runs the full
    `k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
    reproducing the pre-compaction behaviour. Required by
    `_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
    whose training contract needs a real prediction at every slot up to
    `k_max` (mixed per-slot against ground truth) even past a row's own
    `n_sec` — see that function's docstring.

    Under `history="attention"` the history encoding is computed once per
    slot via `Stage2Autoregressive.history_step` (a KV-cache append)
    rather than re-derived by every model call inside that slot — so an ODE
    loop's `steps` substeps, and the separate `predict_type` call when the
    type slice isn't folded into the trunk output, all reuse the SAME `hist`
    tensor for a given `k`. Recomputing per call instead would be merely
    wasteful under markov (its per-call cost is already O(1)) but wrong under
    attention: `AttentionHistory.step` mutates the cache by appending, so
    calling it more than once per slot would double-count that slot's own
    (not-yet-existing) predecessor.

    The free-running history feature stays UNSNAPPED (mirrors the
    established "no snapping" precedent for `particle_type.target =
    "physical"` secondaries feeding their own future conditioning):
    `"physical"` carries the raw (log_mass, charge) forward as-is;
    `"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
    one exception — its history slot must be a probability-simplex-shaped
    vector (that's what `MarkovHistory`/`AttentionHistory` were trained on,
    `_type_repr`'s `F.one_hot` ground truth), so it's the hard one-hot of
    `argmax(logits)`, not the raw logits themselves. Discretizing further,
    into a concrete PDG code, only ever happens once — at secondary-spawn
    time in `giant/rollout.py` — never inside this loop.

    Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
    `sample_secondaries`/`sample_secondaries_wgan`'s (see
    `_decode_stage2_flat`'s docstring); `sec_type` is raw per-slot output in
    all three `particle_type.target` cases (never one-hot-collapsed), so the
    caller decodes it exactly the same way regardless of which decoder
    produced it.
    """
    sec_decoder.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    k_max = sec_decoder.k_max
    type_dim = sec_decoder.type_dim
    objective = build_objective(sec_decoder.generator_kind)
    target = sec_decoder.particle_type_cfg.target
    type_folded = _type_folded(sec_decoder)
    token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM

    sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
    sec_type = torch.zeros(B, k_max, type_dim, device=device)

    use_stop_token = n_sec_pred is None
    if use_stop_token:
        assert getattr(sec_decoder, "stop_head", None) is not None, (
            "sample_secondaries_ar called with n_sec_pred=None on a decoder "
            "with no stop_head — only valid under stage2_model.n_sec.mode = "
            "'stop_token'"
        )
        derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
        # Tracks which rows have already recorded a stop, globally by
        # original batch index — needed even under compaction's own
        # never-revisit guarantee, because `full_length=True` keeps every
        # row in `active_idx` for the whole loop, so a row whose stop logit
        # fires once but flips back below threshold at a later slot (a real
        # possibility for an untrained/lightly-trained stop_head) must not
        # have `derived_n_sec` overwritten by that later, spurious re-fire.
        finished = torch.zeros(B, dtype=torch.bool, device=device)

    # `active_idx`: rows still contributing tokens, indexed into the
    # original batch. Only ever shrinks (never full_length) or stays fixed
    # at arange(B) (full_length) — see the row-compaction docstring section.
    active_idx = torch.arange(B, device=device)
    if not full_length and not use_stop_token:
        active_idx = active_idx[n_sec_pred > 0]

    # Running per-token state, already compacted to `active_idx`.
    prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
    remaining = torch.ones(active_idx.numel(), device=device)
    history_cache = sec_decoder.init_history_cache()

    for k in range(k_max):
        if active_idx.numel() == 0:
            break
        Bc = active_idx.numel()
        cc = cond_cont.index_select(0, active_idx)
        ck = cond_cat.index_select(0, active_idx)
        s1 = stage1_out.index_select(0, active_idx)
        has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
        history_feat = prev_repr.unsqueeze(1)  # (Bc, 1, CONT_SLOT_DIM + type_dim)
        remaining_frac = remaining.unsqueeze(1)  # (Bc, 1)
        slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
        hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)

        if use_stop_token:
            stop_logit = sec_decoder.predict_stop(
                cc,
                ck,
                s1,
                history_feat,
                has_prev,
                remaining_frac,
                slot_idx,
                hist=hist,
            ).squeeze(1)
            if sec_decoder.n_sec_sampling == "sample":
                stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
            else:
                stop_now = stop_logit >= 0.0
            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
            derived_n_sec[active_idx[newly_stopped]] = k
            finished[active_idx[stop_now]] = True
            if not full_length:
                keep = ~stop_now
                active_idx = active_idx[keep]
                cc, ck, s1 = cc[keep], ck[keep], s1[keep]
                has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
                history_feat, hist = history_feat[keep], hist[keep]
                history_cache = sec_decoder.select_history_cache(history_cache, keep)
                prev_repr, remaining = prev_repr[keep], remaining[keep]
                if active_idx.numel() == 0:
                    break
                Bc = active_idx.numel()

        if objective.is_adversarial:
            z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
            token = sec_decoder(
                z,
                cc,
                ck,
                s1,
                history_feat,
                has_prev,
                remaining_frac,
                slot_idx,
                hist=hist,
            )
        else:
            x = torch.randn(Bc, 1, token_dim, device=device)
            dt = 1.0 / steps
            for i in range(steps):
                t = torch.full((Bc, 1), i * dt, device=device)
                v = sec_decoder(
                    x,
                    cc,
                    ck,
                    s1,
                    history_feat,
                    has_prev,
                    remaining_frac,
                    slot_idx,
                    t=t,
                    hist=hist,
                )
                x = x + v * dt
            token = x

        token = token.squeeze(1)  # (Bc, token_dim)
        cont_k = token[:, :CONT_SLOT_DIM]
        if type_folded:
            type_k = token[:, CONT_SLOT_DIM:]
        else:
            type_k = sec_decoder.predict_type(
                cc,
                ck,
                s1,
                history_feat,
                has_prev,
                remaining_frac,
                slot_idx,
                hist=hist,
            ).squeeze(1)

        sec_cont[active_idx, k] = cont_k
        sec_type[active_idx, k] = type_k

        if target == "onehot":
            type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
        else:
            type_for_history = type_k

        stick_fraction = torch.sigmoid(cont_k[:, 0])
        prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
        remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)

        if not full_length and not use_stop_token:
            keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
            active_idx = active_idx[keep2]
            prev_repr, remaining = prev_repr[keep2], remaining[keep2]
            history_cache = sec_decoder.select_history_cache(history_cache, keep2)

    resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
    return sec_cont, sec_type, sec_valid


# ---------------------------------------------------------------------------
# Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
# `predict` command, since both need "given a stage model, produce a
# sample" without hand-picking the sampler themselves (each stage's
# generative objective is independent, read off the model's own
# `generator_kind`, not a caller-supplied `mode` string).
# ---------------------------------------------------------------------------


def sample_stage1(
    stage1_model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    steps: int,
    ddpm_steps: int = 1000,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Dispatches on `stage1_model.generator_kind`."""
    objective = build_objective(stage1_model.generator_kind)
    if objective.is_adversarial:
        return sample_wgan(stage1_model, cond_cont, cond_cat)
    if isinstance(objective, DdpmObjective):
        schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
        return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
    return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)


def sample_stage2(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor | None,
    steps: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Dispatches on `decoder` (one-shot vs autoregressive — the class
    itself, via `isinstance`) and `sec_decoder.generator_kind` (flow/ddpm/
    wgan). DDPM secondaries aren't supported — no `Stage2*` class was ever
    built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
    is the only stage-2 training path that exists for the non-adversarial
    case, so there's nothing to dispatch to here.

    `n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only
    meaningful for the autoregressive path — see `sample_secondaries_ar`'s
    docstring; the one-shot samplers have no per-token stop mechanism to
    derive a count from, so `n_sec_pred` must already be resolved for them.
    """
    if isinstance(sec_decoder, Stage2Autoregressive):
        return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
    assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred"
    if build_objective(sec_decoder.generator_kind).is_adversarial:
        return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
    return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)


def resolve_n_sec(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor | None,
) -> torch.Tensor | None:
    """`n_sec_pred` is already populated when `stage1_model` owns a legacy
    `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
    docstring); otherwise ask stage 2, which owns it by default.

    Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode =
    "stop_token"`) instead of an `n_sec_head` — there is nothing to resolve
    up front in that case, since the count only exists once
    `sample_secondaries_ar` has actually generated (or stopped generating)
    tokens; the caller passes this `None` straight through to `sample_stage2`
    and reads the real count back off its returned `sec_valid`
    (`sec_valid.sum(-1)`) afterwards.

    Raises if neither stage owns any n_sec mechanism at all — the only way
    that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
    rollout-/predict-capable checkpoint.

    `n_sec.mode = "head"` resolves the classifier logits per
    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
    mode via argmax; "sample" draws a real sample from the learned count
    distribution via `torch.multinomial` on the softmax — see gitea #86."""
    if n_sec_pred is not None:
        return n_sec_pred
    if getattr(sec_decoder, "stop_head", None) is not None:
        return None
    if getattr(sec_decoder, "n_sec_head", None) is None:
        raise RuntimeError(
            "checkpoint has no n_sec_head/stop_head on either stage — needs "
            "stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; "
            "'truth' is standalone-evaluation-only"
        )
    logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
    if sec_decoder.n_sec_sampling == "sample":
        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
    return logits.argmax(dim=-1)