lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
5c576fa8f3f88b2f481f8084ee5567237c666762
giant/tests/test_sample.py
T
lars 5c576fa8f3
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 1m2s
Details
CI / Lint (ruff check) (pull_request) Successful in 1m6s
Details
CI / Format (ruff format) (pull_request) Successful in 1m6s
Details
CI / Tests (pull_request) Successful in 2m46s
Details
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
Details
CI / Publish package to Gitea package registry (pull_request) Has been skipped
Details
CI / Update README badges (version, test count) (pull_request) Has been skipped
Details
perf: compact Stage-2 AR inference loop to active rows only 
sample_secondaries_ar ran all k_max=15 slots for every row regardless of
each row's own predicted secondary count, even though the baseline
checkpoint's rollout measured only 0.382 secondaries/step — so ~97% of
stage-2 model calls generated tokens sec_valid then masked away.

Compact the loop to the still-active row set at each slot: drop a row the
moment its n_sec_pred is exhausted (or, under n_sec.mode="stop_token", the
moment its own stop logit fires), so slot k's model calls cost O(active
rows) instead of O(B). Exact — rows are independent given their own
history — verified by comparing the compacted path against a new
full_length=True escape hatch that reproduces the original uncompacted
behavior bit-for-bit under deterministic noise.

full_length=True is 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 regardless of a row's own count, so training behavior is
unchanged.

AttentionHistory's KV cache and MarkovHistory's O(1) state are kept
aligned to the shrinking active set via a new
HistoryEncoder.select_cache / Stage2Autoregressive.select_history_cache.

Also fixes a latent bug the refactor surfaced: derived_n_sec (stop-token
mode) could be overwritten by a later spurious re-fire of the stop logit
on a row that had already stopped; now tracked via an explicit `finished`
mask so only the first stop slot is recorded, matching the documented
contract.

No architecture or checkpoint-format change — every existing v0.3.0
Stage2Autoregressive checkpoint (flow/wgan, markov/attention,
head/stop_token) picks up the speedup automatically on its next
`giant rollout`/`giant predict`, no retraining needed.

Measured (CPU, hidden_dim=512/6 blocks, k_max=15, batch 512, mean
n_sec≈0.38 matching the baseline checkpoint's own rollout): 17.6-22.9x
fewer wall-clock seconds for the AR loop alone (attention/markov history
respectively). Directional only — baseline.toml's GPU inference-cost
comment is updated accordingly, flagged stale pending a real rollout
re-measurement via eval_cost_per_step.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HPt7bVLZYFJe5cG6V7ahqC
2026-09-03 17:56:59 +02:00

480 lines
20 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 (
resolve_n_sec,
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",
n_sec_sampling: str = "greedy",
) -> 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,
n_sec_sampling=n_sec_sampling,
).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,
n_sec_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,
n_sec_sampling=n_sec_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("n_sec_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_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", n_sec_sampling=n_sec_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("n_sec_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_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", n_sec_sampling=n_sec_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)
# ── Row compaction: full_length=False (default, inference) must agree with
# full_length=True (the pre-compaction behaviour, still exercised by
# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────
def _zero_randn(*size, **kwargs):
"""Drop-in replacement for `torch.randn` that returns zeros of the same
shape — makes the ODE/WGAN noise deterministic so a compacted run and a
full_length run can be compared row-for-row regardless of how many
`torch.randn` calls each makes (compaction changes the batch size, and
therefore the RNG stream position, at every slot)."""
device = kwargs.get("device")
dtype = kwargs.get("dtype")
return torch.zeros(*size, device=device, dtype=dtype)
@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
B, k_max, emb_dim = 4, 5, 6
decoder = _stage2_ar("physical", 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.tensor([0, 1, 3, k_max])
monkeypatch.setattr(torch, "randn", _zero_randn)
sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
)
sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
)
assert torch.equal(sec_valid_c, sec_valid_f)
assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
"""`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
(no `torch.rand` draw), so only `torch.randn` needs zeroing."""
B, k_max = 6, 5
decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
monkeypatch.setattr(torch, "randn", _zero_randn)
sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
)
sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
)
assert torch.equal(sec_valid_c, sec_valid_f)
assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
"""A row with n_sec_pred == 0 would be dropped from the active set at
slot 0 under compaction (full_length=False) — full_length=True must
still run the model for it at every slot (only masked by sec_valid at
the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
B, k_max, emb_dim = 3, 4, 6
decoder = _stage2_ar("physical", "flow", 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, 0, 0])
sec_cont, sec_type, sec_valid = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
)
assert not sec_valid.any()
# every slot still ran the model (not left at the zero-init default) —
# a real flow ODE output from randn-initialized noise is essentially
# never exactly zero.
assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))
# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
"""Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
returns `bias` (broadcast over the batch) as logits regardless of
conditioning — mirrors `_force_stop_head_logit`."""
assert decoder.n_sec_head is not None
last_linear = decoder.n_sec_head[-1]
with torch.no_grad():
last_linear.weight.zero_()
last_linear.bias.copy_(bias)
@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
"""A logit vector overwhelmingly favoring one class gives the same
answer under both policies — greedy because it's the argmax, sample
because softmax puts ~all mass on it."""
B, k_max = 8, 5
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
bias = torch.full((k_max + 1,), -50.0)
bias[2] = 50.0
_force_n_sec_head_bias(decoder, bias)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
assert n_sec is not None
assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
B, k_max = 32, 5
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
_force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
assert n_sec is not None
assert n_sec.unique().numel() == 1
def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
"""Under a flat logit vector, a categorical draw across a large batch
should hit more than one class — the whole point of gitea #86: greedy
always collapses to one, sample should not."""
torch.manual_seed(0)
B, k_max = 256, 5
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
_force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
assert n_sec is not None
assert n_sec.unique().numel() > 1
Reference in New Issue View Git Blame Copy Permalink