lars/giant
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
8290e350b81e8c544215b648ecfd2fb5a93fb046
giant/tests/test_sample.py
T
lars c1c4957e2f Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40) 
Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:

- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
  P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
  (n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
  stop head instead of a pre-resolved count: each row stops the first slot
  its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
  default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
  batch loop breaks once every row has stopped, so cost scales with the
  realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
  (the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
  rollout.py/cli.py/validate.py now derive the realized count from
  sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
  rather than resolving it up front — a no-op reordering under every other
  n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
  the per-slot target/mask (one slot wider than the existing token-content
  sec_mask, since the stop slot itself needs supervision) and
  StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
  exactly like _n_sec_loss gates on n_sec_head. Wired into both the
  flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
  stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
  head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
  and n_sec.owner="stage2") instead of always rejecting it.

Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 10:44:14 +02:00

339 lines
13 KiB
Python
Raw Blame History

"""Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop
(`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage
for the one-shot samplers."""
import pytest
import torch
from giant.config import ConditioningAxisConfig, ParticleTypeConfig
from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM
from giant.model.network import (
Stage1Model,
Stage2Autoregressive,
Stage2OneShot,
stage2_trunk_sec_dim,
)
from giant.sample import (
sample_flow,
sample_secondaries,
sample_secondaries_ar,
sample_secondaries_wgan,
sample_wgan,
)
_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]:
cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
return cfg, cfg
def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
return cond_cont, cond_cat
def _conditioning_for(target: str) -> str:
# target="embedding" regresses against the conditioning's own embedding
# table — only meaningful when the
# conditioning axis is itself "embedding".
return "embedding" if target == "embedding" else "physical"
def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
particle_type_cfg = ParticleTypeConfig(target=target)
# build_models (giant/model/network.py) computes sec_dim this same way
# before constructing Stage2OneShot — its own default (SEC_DIM, the
# "physical" width) is only correct for target="physical".
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, K_MAX, emb_dim)
return Stage2OneShot(
pdg_vocab=pdg,
mat_vocab=mat,
particle_cfg=particle_cfg,
material_cfg=material_cfg,
hidden_dim=32,
n_res_blocks=2,
generator=generator,
time_dim=16,
noise_dim=8,
sec_dim=sec_dim,
particle_type_cfg=particle_type_cfg,
).eval()
def _stage2_ar(
target: str,
generator: str,
emb_dim: int = 6,
pdg: int = 3,
mat: int = 2,
k_max: int = 5,
history: str = "markov",
) -> Stage2Autoregressive:
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
return Stage2Autoregressive(
pdg_vocab=pdg,
mat_vocab=mat,
particle_cfg=particle_cfg,
material_cfg=material_cfg,
hidden_dim=32,
n_res_blocks=2,
generator=generator,
time_dim=16,
noise_dim=8,
k_max=k_max,
particle_type_cfg=ParticleTypeConfig(target=target),
history=history,
attn_n_heads=2,
attn_n_layers=1,
).eval()
def _expected_type_dim(target: str, emb_dim: int) -> int:
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
def _stage2_ar_stop_token(
target: str,
generator: str,
stop_sampling: str = "greedy",
emb_dim: int = 6,
pdg: int = 3,
mat: int = 2,
k_max: int = 5,
) -> Stage2Autoregressive:
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
return Stage2Autoregressive(
pdg_vocab=pdg,
mat_vocab=mat,
particle_cfg=particle_cfg,
material_cfg=material_cfg,
hidden_dim=32,
n_res_blocks=2,
generator=generator,
time_dim=16,
noise_dim=8,
k_max=k_max,
particle_type_cfg=ParticleTypeConfig(target=target),
build_n_sec_head=False,
build_stop_head=True,
stop_sampling=stop_sampling,
).eval()
def _force_stop_head_logit(decoder: Stage2Autoregressive, logit: float) -> None:
"""Zeroes stop_head's weights and pins its bias, so predict_stop returns
`logit` for every row/slot regardless of conditioning — makes the AR
loop's stop decision deterministic for testing."""
assert decoder.stop_head is not None
last_linear = decoder.stop_head[-1]
with torch.no_grad():
last_linear.weight.zero_()
last_linear.bias.fill_(logit)
# ── Stage-1 n_sec ownership ──────────────────────────────────────────────────
def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head():
model = Stage1Model(
pdg_vocab=3,
mat_vocab=2,
particle_cfg=_PHYS_CFG,
material_cfg=_PHYS_CFG,
hidden_dim=16,
n_res_blocks=1,
)
cond_cont, cond_cat = _cond(4)
sample, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
assert sample.shape == (4, X_DIM)
assert n_sec is None
def test_sample_wgan_returns_none_n_sec_when_stage1_owns_no_head():
model = Stage1Model(
pdg_vocab=3,
mat_vocab=2,
particle_cfg=_PHYS_CFG,
material_cfg=_PHYS_CFG,
hidden_dim=16,
n_res_blocks=1,
generator="wgan",
noise_dim=8,
)
cond_cont, cond_cat = _cond(4)
sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
assert sample.shape == (4, X_DIM)
assert n_sec is None
def test_sample_flow_returns_n_sec_for_legacy_stage1():
model = Stage1Model(
pdg_vocab=3,
mat_vocab=2,
particle_cfg=_PHYS_CFG,
material_cfg=_PHYS_CFG,
hidden_dim=16,
n_res_blocks=1,
n_sec_head_k_max=K_MAX,
)
cond_cont, cond_cat = _cond(5)
_, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
assert n_sec is not None and n_sec.shape == (5,)
# ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_flow_shapes_by_target(target):
B, emb_dim = 5, 6
decoder = _stage2_oneshot(target, "flow", emb_dim=emb_dim)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, K_MAX)
assert torch.isfinite(sec_cont).all()
assert torch.isfinite(sec_type).all()
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_wgan_shapes_by_target(target):
B, emb_dim = 5, 6
decoder = _stage2_oneshot(target, "wgan", emb_dim=emb_dim)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, K_MAX)
# ── Stage2Autoregressive ─────────────────────────────────────────────────────
@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("generator", ["flow", "wgan"])
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_shapes(target, generator, history):
B, k_max, emb_dim = 4, 5, 6
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, k_max + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, k_max)
assert torch.isfinite(sec_cont).all()
assert torch.isfinite(sec_type).all()
@pytest.mark.parametrize("generator", ["flow", "wgan"])
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
B, k_max, emb_dim = 3, 5, 6
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 2, k_max])
_, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
assert not sec_valid[i, n:].any()
def test_sample_secondaries_ar_first_slot_has_no_history():
"""Slot 0 always has has_prev=False internally — nothing to assert on
the public API directly, but a k_max=1 run should not crash on the
"previous token" path at all (has_prev never true)."""
B, emb_dim = 3, 6
decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=1)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 1, 1])
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
assert sec_valid.tolist() == [[False], [True], [True]]
# ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
"""A stop_head pinned to a large positive logit fires at slot 0 for
every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
generating any token."""
B, k_max = 4, 5
decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
_force_stop_head_logit(decoder, 50.0)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
assert sec_valid.shape == (B, k_max)
assert not sec_valid.any()
@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
"""A stop_head pinned to a large negative logit never fires under either
policy, so every row is capped at k_max (the safety cap, not a modeling
ceiling)."""
B, k_max = 4, 5
decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
_force_stop_head_logit(decoder, -50.0)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
assert sec_valid.all()
@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_stop_token_valid_mask_is_always_a_prefix(generator):
"""Without forcing the stop head, per-row stop timing varies — but
sec_valid must always be a contiguous prefix (slot k valid implies every
slot < k is also valid), matching the "head"/"truth" contract."""
B, k_max = 6, 5
decoder = _stage2_ar_stop_token("physical", generator, k_max=k_max)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
_, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
n = sec_valid.sum(dim=-1)
expected = torch.arange(k_max).unsqueeze(0) < n.unsqueeze(1)
assert torch.equal(sec_valid, expected)
def test_sample_secondaries_ar_stop_token_explicit_n_sec_pred_ignores_stop_head():
"""The scheduled-sampling training contract: passing n_sec_pred
explicitly (as _assemble_stage2_ar_inputs_scheduled's self-sample call
does, with ground-truth n_sec) must run the full k_max loop and mask by
the given count, even though the decoder owns a stop_head that would
otherwise stop early."""
B, k_max = 3, 5
decoder = _stage2_ar_stop_token("physical", "flow", k_max=k_max)
_force_stop_head_logit(decoder, 50.0) # would stop immediately if consulted
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 2, k_max])
_, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
assert not sec_valid[i, n:].any()
def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
decoder = _stage2_ar("physical", "flow", k_max=5) # head mode: no stop_head
cond_cont, cond_cat = _cond(3)
stage1_out = torch.randn(3, X_DIM)
with pytest.raises(AssertionError):
sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
Reference in New Issue View Git Blame Copy Permalink