lars
200c6d243b
AttentionHistory (giant/model/network.py) adds causal self-attention over the emitted-secondary prefix as the alternative to MarkovHistory, with a parallel forward() for training and an init_cache()/step() KV-cache path for sample.py's per-slot AR inference loop, wired into Stage2Autoregressive via history="attention". giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled, mixing ground-truth history with a detached sample_secondaries_ar self-sample per slot so teacher_forcing="scheduled"/"never" close the train/inference gap teacher_forcing="always" always avoided; wired into both stage-2 AR trainers. config.py's validate_config no longer rejects these two previously unimplemented schema values. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
631 lines
21 KiB
Python
631 lines
21 KiB
Python
import pytest
|
|
import torch
|
|
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
|
|
from giant.model.network import (
|
|
AttentionHistory,
|
|
ConditionEncoder,
|
|
MarkovHistory,
|
|
SinusoidalEmbedding,
|
|
Stage1Model,
|
|
Stage2Autoregressive,
|
|
Stage2OneShot,
|
|
cat_col_layout,
|
|
stage2_trunk_sec_dim,
|
|
stage2_type_dim,
|
|
)
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
|
|
ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
|
|
|
|
|
|
def test_sinusoidal_embedding_shape():
|
|
emb = SinusoidalEmbedding(64)
|
|
t = torch.rand(16)
|
|
assert emb(t).shape == (16, 64)
|
|
|
|
|
|
def test_sinusoidal_embedding_batch_1():
|
|
emb = SinusoidalEmbedding(32)
|
|
t = torch.tensor([0.5])
|
|
assert emb(t).shape == (1, 32)
|
|
|
|
|
|
def test_stage1_model_output_shape():
|
|
B = 8
|
|
model = Stage1Model(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
x_t = torch.randn(B, 9)
|
|
t = torch.rand(B)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.randint(0, 5, (B,)),
|
|
torch.randint(0, 3, (B,)),
|
|
],
|
|
dim=1,
|
|
)
|
|
out = model(x_t, cond_cont, cond_cat, t=t)
|
|
assert out.shape == (B, 9)
|
|
|
|
|
|
def test_stage1_model_gradients_flow():
|
|
B = 4
|
|
model = Stage1Model(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=32,
|
|
n_res_blocks=2,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
x_t = torch.randn(B, 9)
|
|
t = torch.rand(B)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
# Both paths must be exercised to get gradients through all parameters.
|
|
flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
|
|
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
|
|
(flow_loss + nsec_loss).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
def test_stage1_model_no_n_sec_head_by_default():
|
|
"""Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
|
|
decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2."""
|
|
model = Stage1Model(
|
|
pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
|
|
)
|
|
assert model.n_sec_head is None
|
|
|
|
|
|
# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
|
|
|
|
|
|
def test_cat_col_layout_neither_onehot():
|
|
assert cat_col_layout("physical", "embedding") == (None, None)
|
|
|
|
|
|
def test_cat_col_layout_particle_only():
|
|
assert cat_col_layout("onehot", "physical") == (2, None)
|
|
|
|
|
|
def test_cat_col_layout_material_only():
|
|
assert cat_col_layout("physical", "onehot") == (None, 2)
|
|
|
|
|
|
def test_cat_col_layout_both_onehot_particle_then_material():
|
|
assert cat_col_layout("onehot", "onehot") == (2, 3)
|
|
|
|
|
|
def test_stage2_type_dim_physical_is_particle_phys_dim():
|
|
assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
|
|
|
|
|
|
def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
|
|
assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
|
|
assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
|
|
|
|
|
|
def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
|
|
k_max = 15
|
|
assert (
|
|
stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
|
|
== k_max * SEC_SLOT_DIM
|
|
)
|
|
assert (
|
|
stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
|
|
== k_max * SEC_SLOT_DIM
|
|
)
|
|
|
|
|
|
def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
|
|
k_max = 15
|
|
assert stage2_trunk_sec_dim(
|
|
{"target": "onehot"}, "wgan", k_max, emb_dim=16
|
|
) == k_max * (CONT_SLOT_DIM + 16)
|
|
|
|
|
|
def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
|
|
k_max = 15
|
|
assert (
|
|
stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
|
|
== k_max * CONT_SLOT_DIM
|
|
)
|
|
|
|
|
|
# --- ConditionEncoder onehot mode -------------------------------------------
|
|
|
|
|
|
def test_condition_encoder_onehot_forward_shape_and_gradients():
|
|
B = 8
|
|
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
|
material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
|
|
enc = ConditionEncoder(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=ONEHOT_PARTICLE_CFG,
|
|
material_cfg=ONEHOT_MATERIAL_CFG,
|
|
out_dim=32,
|
|
)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.randint(0, 5, (B,)),
|
|
torch.randint(0, 3, (B,)),
|
|
torch.randint(0, particle_emb_dim, (B,)),
|
|
torch.randint(0, material_emb_dim, (B,)),
|
|
],
|
|
dim=1,
|
|
)
|
|
out = enc(cond_cont, cond_cat)
|
|
assert out.shape == (B, 32)
|
|
# onehot itself is unlearned, but the fusion MLP downstream still has
|
|
# gradients — the encoder as a whole must still be trainable.
|
|
out.sum().backward()
|
|
assert enc.mlp[0].weight.grad is not None
|
|
|
|
|
|
def test_condition_encoder_onehot_is_a_true_one_hot_vector():
|
|
"""The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
|
|
verify the concatenated input segment really is one-hot, not e.g. an
|
|
accidentally-learned embedding."""
|
|
B = 4
|
|
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
|
enc = ConditionEncoder(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=ONEHOT_PARTICLE_CFG,
|
|
material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
|
|
out_dim=16,
|
|
)
|
|
cond_cont = torch.zeros(B, COND_DIM)
|
|
idx = torch.tensor([0, 1, 2, 5])
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.zeros(B, dtype=torch.long),
|
|
torch.zeros(B, dtype=torch.long),
|
|
idx.clamp(max=particle_emb_dim - 1),
|
|
],
|
|
dim=1,
|
|
)
|
|
pdg_e = enc._particle_embed(cond_cont, cond_cat)
|
|
assert pdg_e.shape == (B, particle_emb_dim)
|
|
assert torch.all(pdg_e.sum(dim=-1) == 1.0)
|
|
|
|
|
|
# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) --
|
|
|
|
|
|
def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
|
|
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
|
if target != "physical":
|
|
particle_cfg = dict(particle_cfg)
|
|
if target == "embedding":
|
|
particle_cfg["type"] = "embedding"
|
|
k_max = 5
|
|
sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
|
|
return Stage2OneShot(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=16,
|
|
n_res_blocks=1,
|
|
cond_out_dim=16,
|
|
context_dim=8,
|
|
sec_dim=sec_dim,
|
|
generator=generator,
|
|
k_max=k_max,
|
|
particle_type_cfg={"target": target, "lambda": 1.0},
|
|
)
|
|
|
|
|
|
def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
|
|
assert _build_stage2("physical", "flow").type_head is None
|
|
assert _build_stage2("physical", "wgan").type_head is None
|
|
|
|
|
|
def test_stage2_oneshot_onehot_flow_has_type_head():
|
|
model = _build_stage2("onehot", "flow")
|
|
assert model.type_head is not None
|
|
|
|
|
|
def test_stage2_oneshot_onehot_wgan_has_no_type_head():
|
|
"""Under wgan the type slice is folded into forward()'s own output and
|
|
relaxed via ST-Gumbel by the trainer — no separate head needed."""
|
|
model = _build_stage2("onehot", "wgan")
|
|
assert model.type_head is None
|
|
|
|
|
|
def test_stage2_oneshot_embedding_flow_has_type_head():
|
|
model = _build_stage2("embedding", "flow")
|
|
assert model.type_head is not None
|
|
|
|
|
|
def test_stage2_oneshot_predict_type_shape():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
out = model.predict_type(cond_cont, cond_cat, stage1_out)
|
|
assert out.shape == (B, k_max, emb_dim)
|
|
|
|
|
|
def test_stage2_oneshot_predict_type_raises_when_no_type_head():
|
|
model = _build_stage2("physical", "flow")
|
|
cond_cont = torch.randn(2, COND_DIM)
|
|
cond_cat = torch.zeros(2, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(2, 9)
|
|
try:
|
|
model.predict_type(cond_cont, cond_cat, stage1_out)
|
|
raise AssertionError("expected RuntimeError")
|
|
except RuntimeError:
|
|
pass
|
|
|
|
|
|
def test_stage2_oneshot_forward_shape_onehot_wgan():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
z = torch.randn(B, model.noise_dim)
|
|
out = model(z, cond_cont, cond_cat, stage1_out)
|
|
assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))
|
|
|
|
|
|
def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
|
|
t = torch.rand(B)
|
|
out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
|
assert out.shape == (B, k_max * CONT_SLOT_DIM)
|
|
|
|
|
|
# --- MarkovHistory (docs/v0.3.0-design.md §6.2) -----------------------------
|
|
|
|
|
|
def test_markov_history_shape():
|
|
hist = MarkovHistory(in_dim=7, out_dim=12)
|
|
B, K = 3, 5
|
|
feat = torch.randn(B, K, 7)
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
out = hist(feat, has_prev)
|
|
assert out.shape == (B, K, 12)
|
|
|
|
|
|
def test_markov_history_uses_start_vector_when_no_prev():
|
|
"""Slot 0's own raw feature must be ignored — a learned start vector is
|
|
substituted there instead (a reasonable default not specified by the
|
|
design doc, see Stage2Autoregressive's docstring)."""
|
|
hist = MarkovHistory(in_dim=4, out_dim=6)
|
|
B, K = 2, 3
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
feat_a = torch.randn(B, K, 4)
|
|
feat_b = feat_a.clone()
|
|
feat_b[:, 0] = torch.randn(B, 4) * 100
|
|
out_a = hist(feat_a, has_prev)
|
|
out_b = hist(feat_b, has_prev)
|
|
assert torch.allclose(out_a[:, 0], out_b[:, 0])
|
|
assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
|
|
|
|
|
|
# --- AttentionHistory (docs/v0.3.0-design.md §6.2, v0.3.0 step 7) ----------
|
|
|
|
|
|
def test_attention_history_shape():
|
|
hist = AttentionHistory(in_dim=7, out_dim=12, n_heads=2, n_layers=2)
|
|
B, K = 3, 5
|
|
feat = torch.randn(B, K, 7)
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
out = hist(feat, has_prev)
|
|
assert out.shape == (B, K, 12)
|
|
|
|
|
|
def test_attention_history_uses_start_vector_when_no_prev():
|
|
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=1)
|
|
B, K = 2, 3
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
feat_a = torch.randn(B, K, 4)
|
|
feat_b = feat_a.clone()
|
|
feat_b[:, 0] = torch.randn(B, 4) * 100
|
|
out_a = hist(feat_a, has_prev)
|
|
out_b = hist(feat_b, has_prev)
|
|
assert torch.allclose(out_a[:, 0], out_b[:, 0], atol=1e-5)
|
|
|
|
|
|
def test_attention_history_is_causal():
|
|
"""Position i's output must not depend on feat at positions > i — unlike
|
|
MarkovHistory (which only ever looks at position i itself, already
|
|
trivially "causal"), this is AttentionHistory's actual contribution:
|
|
seeing the full prefix 0..i-1, never anything later."""
|
|
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
|
|
hist.eval()
|
|
B, K = 2, 5
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
feat_a = torch.randn(B, K, 4)
|
|
feat_b = feat_a.clone()
|
|
feat_b[:, 3:] = torch.randn(B, K - 3, 4) * 100
|
|
with torch.no_grad():
|
|
out_a = hist(feat_a, has_prev)
|
|
out_b = hist(feat_b, has_prev)
|
|
assert torch.allclose(out_a[:, :3], out_b[:, :3], atol=1e-5)
|
|
|
|
|
|
def test_attention_history_step_matches_forward():
|
|
"""The incremental KV-cache path (`init_cache`/`step`,
|
|
`giant/sample.py`'s AR loop) must reproduce `forward`'s parallel-pass
|
|
output exactly, one position at a time."""
|
|
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
|
|
hist.eval()
|
|
B, K = 3, 6
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
feat = torch.randn(B, K, 4)
|
|
with torch.no_grad():
|
|
expected = hist(feat, has_prev)
|
|
|
|
cache = hist.init_cache()
|
|
outs = []
|
|
for k in range(K):
|
|
out_k, cache = hist.step(feat[:, k : k + 1], has_prev[:, k : k + 1], cache)
|
|
outs.append(out_k)
|
|
stepped = torch.cat(outs, dim=1)
|
|
|
|
assert torch.allclose(stepped, expected, atol=1e-5)
|
|
|
|
|
|
# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) ---------
|
|
|
|
|
|
def _build_stage2_ar(
|
|
target: str,
|
|
generator: str,
|
|
emb_dim: int = 6,
|
|
k_max: int = 5,
|
|
history: str = "markov",
|
|
) -> Stage2Autoregressive:
|
|
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
|
if target == "embedding":
|
|
particle_cfg = dict(particle_cfg)
|
|
particle_cfg["type"] = "embedding"
|
|
return Stage2Autoregressive(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=16,
|
|
n_res_blocks=1,
|
|
cond_out_dim=16,
|
|
context_dim=8,
|
|
generator=generator,
|
|
k_max=k_max,
|
|
particle_type_cfg={"target": target, "lambda": 1.0},
|
|
history=history,
|
|
)
|
|
|
|
|
|
def _ar_inputs(B: int, K: int, hist_dim: int):
|
|
history_feat = torch.randn(B, K, hist_dim)
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
remaining_frac = torch.rand(B, K)
|
|
slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1)
|
|
return history_feat, has_prev, remaining_frac, slot_idx
|
|
|
|
|
|
def test_stage2_autoregressive_history_invalid_raises():
|
|
with pytest.raises(ValueError):
|
|
_build_stage2_ar("onehot", "wgan", history="bogus")
|
|
|
|
|
|
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
|
@pytest.mark.parametrize("generator", ["wgan", "flow"])
|
|
@pytest.mark.parametrize("history", ["markov", "attention"])
|
|
def test_stage2_autoregressive_forward_shape(target, generator, history):
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar(
|
|
target, generator, emb_dim=emb_dim, k_max=K, history=history
|
|
)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
type_dim = stage2_type_dim({"target": target}, emb_dim)
|
|
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
|
|
B, K, CONT_SLOT_DIM + type_dim
|
|
)
|
|
token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
|
|
if generator == "wgan":
|
|
x_t = torch.randn(B, K, model.noise_dim)
|
|
t = None
|
|
else:
|
|
x_t = torch.randn(B, K, token_dim)
|
|
t = torch.rand(B, K)
|
|
out = model(
|
|
x_t,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
t=t,
|
|
)
|
|
assert out.shape == (B, K, token_dim)
|
|
|
|
|
|
def test_stage2_autoregressive_predict_n_sec_shape():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=k_max)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
logits = model.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
|
assert logits.shape == (B, k_max + 1)
|
|
|
|
|
|
def test_stage2_autoregressive_predict_type_shape():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
|
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
|
|
B, K, CONT_SLOT_DIM + type_dim
|
|
)
|
|
out = model.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
)
|
|
assert out.shape == (B, K, emb_dim)
|
|
|
|
|
|
@pytest.mark.parametrize("target,generator", [("physical", "flow"), ("onehot", "wgan")])
|
|
def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, generator):
|
|
B, K, emb_dim = 2, 5, 6
|
|
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
type_dim = stage2_type_dim({"target": target}, emb_dim)
|
|
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
|
|
B, K, CONT_SLOT_DIM + type_dim
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
model.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
)
|
|
|
|
|
|
def test_stage2_autoregressive_gradients_flow_wgan_onehot():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "wgan", emb_dim=emb_dim, k_max=K)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
|
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
|
|
B, K, CONT_SLOT_DIM + type_dim
|
|
)
|
|
z = torch.randn(B, K, model.noise_dim)
|
|
gen_out = model(
|
|
z,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
).sum()
|
|
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
|
(gen_out + nsec_out).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
def test_stage2_autoregressive_gradients_flow_onehot():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
|
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
|
|
B, K, CONT_SLOT_DIM + type_dim
|
|
)
|
|
token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
|
|
x_t = torch.randn(B, K, token_dim)
|
|
t = torch.rand(B, K)
|
|
flow_out = model(
|
|
x_t,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
t=t,
|
|
).sum()
|
|
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
|
type_out = model.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
).sum()
|
|
(flow_out + nsec_out + type_out).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
|
|
"""`init_history_cache`/`history_step` (the incremental path
|
|
`giant/sample.py`'s AR loop drives, one slot per call) must reproduce
|
|
exactly what one parallel `self.history_encoder(history_feat, has_prev)`
|
|
call over the whole shifted sequence would give at each position — the
|
|
KV-cache correctness guarantee, exercised through `Stage2Autoregressive`
|
|
itself rather than `AttentionHistory` in isolation
|
|
(`test_attention_history_step_matches_forward` covers that lower layer)."""
|
|
B, K, emb_dim = 3, 6, 6
|
|
model = _build_stage2_ar(
|
|
"physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
|
|
)
|
|
model.eval()
|
|
type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
|
|
hist_in_dim = CONT_SLOT_DIM + type_dim
|
|
own_feat = torch.randn(B, K, hist_in_dim) # token i's own raw feature
|
|
has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
history_feat = torch.cat(
|
|
[torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
|
|
)
|
|
|
|
with torch.no_grad():
|
|
expected = model.history_encoder(history_feat, has_prev_full)
|
|
|
|
cache = model.init_history_cache()
|
|
outs = []
|
|
prev = torch.zeros(B, 1, hist_in_dim)
|
|
for k in range(K):
|
|
has_prev_k = torch.full((B, 1), k >= 1, dtype=torch.bool)
|
|
hist_k, cache = model.history_step(prev, has_prev_k, cache)
|
|
outs.append(hist_k)
|
|
prev = own_feat[:, k : k + 1]
|
|
stepped = torch.cat(outs, dim=1)
|
|
|
|
assert torch.allclose(stepped, expected, atol=1e-5)
|
|
|
|
|
|
def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
|
|
model = _build_stage2_ar("physical", "wgan", history="markov")
|
|
assert model.init_history_cache() is None
|