lars
593c5f4d34
CI / Format (ruff format) (push) Successful in 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 34s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Type check (ty) (pull_request) Successful in 46s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Tests (pull_request) Successful in 4m5s
CI / Tests (push) Successful in 4m6s
The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2)
-> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in
giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/
.type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2`
ratio and fixed 2-layer depth were undocumented magic numbers, and both
n_sec accuracy and secondary-species accuracy are known weak spots that
were untunable independently of the trunk they hang off.
Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to
giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old
hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth)
dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and
`stage2_model.heads.{n_sec,type}` — split per head type (not one shared
block per stage) since n_sec and species prediction are called out as
separate weak spots that may want independent capacity. Defaults
(hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture
bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint
is unaffected; no changes were needed to migrate_config or the legacy
migration surfaces. No new CLI flags, matching how other nested
sub-config (router.*, trunk.*) is set via config.toml rather than
per-field flags.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1006 lines
36 KiB
Python
1006 lines
36 KiB
Python
import copy
|
|
|
|
import pytest
|
|
import torch
|
|
from giant import config as gconfig
|
|
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
|
|
from giant.model.network import (
|
|
HISTORY_REGISTRY,
|
|
AttentionHistory,
|
|
ConditionEncoder,
|
|
HistoryEncoder,
|
|
MarkovHistory,
|
|
SinusoidalEmbedding,
|
|
Stage1Model,
|
|
Stage2Autoregressive,
|
|
Stage2OneShot,
|
|
build_critics,
|
|
build_history,
|
|
build_models,
|
|
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 —
|
|
it moves 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
|
|
|
|
|
|
def test_stage1_model_n_sec_head_default_cfg_matches_pre_gitea_36_shape():
|
|
"""No n_sec_head_cfg given must reproduce the old hardcoded
|
|
hidden_dim // 2, one-hidden-layer architecture exactly (gitea #36)."""
|
|
model = Stage1Model(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=40,
|
|
cond_out_dim=12,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
assert model.n_sec_head is not None
|
|
assert len(model.n_sec_head) == 3
|
|
assert model.n_sec_head[0].in_features == 12
|
|
assert model.n_sec_head[0].out_features == 20 # hidden_dim // 2
|
|
assert model.n_sec_head[2].out_features == 16 # k_max + 1
|
|
|
|
|
|
def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth():
|
|
model = Stage1Model(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=32,
|
|
cond_out_dim=16,
|
|
n_sec_head_k_max=15,
|
|
n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
|
|
)
|
|
assert model.n_sec_head is not None
|
|
assert len(model.n_sec_head) == 1
|
|
assert model.n_sec_head[0].in_features == 16
|
|
assert model.n_sec_head[0].out_features == 16
|
|
|
|
|
|
# --- 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 --------------------------------
|
|
|
|
|
|
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_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
|
|
"""gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable."""
|
|
k_max, emb_dim = 5, 6
|
|
particle_cfg = {"type": "onehot", "emb_dim": emb_dim, "n_layers": 1}
|
|
sec_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim)
|
|
model = Stage2OneShot(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=40,
|
|
n_res_blocks=1,
|
|
cond_out_dim=12,
|
|
context_dim=8,
|
|
sec_dim=sec_dim,
|
|
generator="flow",
|
|
k_max=k_max,
|
|
particle_type_cfg={"target": "onehot", "lambda": 1.0},
|
|
n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
|
|
type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
|
|
)
|
|
assert model.n_sec_head is not None
|
|
assert len(model.n_sec_head) == 1
|
|
assert model.n_sec_head[0].in_features == 12
|
|
assert model.n_sec_head[0].out_features == k_max + 1
|
|
|
|
assert model.type_head is not None
|
|
assert len(model.type_head) == 3
|
|
assert model.type_head[0].out_features == 30 # round(40 * 0.75)
|
|
assert model.type_head[2].out_features == k_max * emb_dim
|
|
|
|
|
|
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)
|
|
|
|
|
|
def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
|
|
"""gitea #29: stage2_model.particle_type.n_classes, not
|
|
conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
|
|
explicitly set — the two used to be silently the same number."""
|
|
k_max = 5
|
|
particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
|
|
particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
|
|
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
|
|
model = 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="flow",
|
|
k_max=k_max,
|
|
particle_type_cfg=particle_type_cfg,
|
|
)
|
|
assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6
|
|
assert model.type_head is not None
|
|
assert model.type_head[-1].out_features == k_max * 20
|
|
|
|
|
|
# --- MarkovHistory -----------------------------------------------------------
|
|
|
|
|
|
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, 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 (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)
|
|
|
|
|
|
# --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
|
|
|
|
|
|
def test_history_registry_has_exactly_the_two_known_histories():
|
|
assert set(HISTORY_REGISTRY) == {"markov", "attention"}
|
|
|
|
|
|
def test_build_history_returns_correct_concrete_type():
|
|
assert isinstance(build_history("markov", 4, 6), MarkovHistory)
|
|
assert isinstance(build_history("attention", 4, 8), AttentionHistory)
|
|
|
|
|
|
def test_build_history_unknown_name_raises():
|
|
with pytest.raises(ValueError):
|
|
build_history("bogus", 4, 6)
|
|
|
|
|
|
def test_build_history_filters_kwargs_by_signature():
|
|
"""Attention-only kwargs (n_heads/n_layers) must be silently dropped when
|
|
building a MarkovHistory, matching build_router's documented behavior for
|
|
per-type hyperparameters coexisting in one config."""
|
|
hist = build_history("markov", 4, 6, n_heads=2, n_layers=1)
|
|
assert isinstance(hist, MarkovHistory)
|
|
|
|
|
|
def test_history_encoder_base_default_init_cache_and_step():
|
|
"""A HistoryEncoder subclass implementing only forward() must still get
|
|
working O(1) init_cache/step defaults from the base class."""
|
|
|
|
class _StubHistory(HistoryEncoder):
|
|
def forward(self, feat, has_prev):
|
|
return feat * 2
|
|
|
|
hist = _StubHistory()
|
|
assert hist.init_cache() is None
|
|
feat = torch.randn(2, 1, 4)
|
|
has_prev = torch.ones(2, 1, dtype=torch.bool)
|
|
out, cache = hist.step(feat, has_prev, "unused-cache")
|
|
assert torch.equal(out, hist.forward(feat, has_prev))
|
|
assert cache == "unused-cache"
|
|
|
|
|
|
# --- Stage2Autoregressive (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")
|
|
|
|
|
|
def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
|
|
"""gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
|
|
of this test for the full rationale."""
|
|
particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
|
|
particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
|
|
model = 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="flow",
|
|
k_max=5,
|
|
particle_type_cfg=particle_type_cfg,
|
|
)
|
|
assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6
|
|
assert model.type_head is not None
|
|
assert model.type_head[-1].out_features == 20
|
|
|
|
|
|
def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
|
|
"""gitea #36, Stage2Autoregressive side — see the Stage2OneShot version
|
|
of this test for the full rationale."""
|
|
particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
|
|
particle_type_cfg = {"target": "onehot", "lambda": 1.0}
|
|
model = Stage2Autoregressive(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=40,
|
|
n_res_blocks=1,
|
|
cond_out_dim=12,
|
|
context_dim=8,
|
|
generator="flow",
|
|
k_max=5,
|
|
particle_type_cfg=particle_type_cfg,
|
|
n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
|
|
type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
|
|
)
|
|
assert model.n_sec_head is not None
|
|
assert len(model.n_sec_head) == 1
|
|
assert model.n_sec_head[0].in_features == 12
|
|
assert model.n_sec_head[0].out_features == 6 # k_max + 1
|
|
|
|
assert model.type_head is not None
|
|
assert len(model.type_head) == 3
|
|
assert model.type_head[0].out_features == 30 # round(40 * 0.75)
|
|
assert model.type_head[2].out_features == model.type_dim
|
|
|
|
|
|
@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
|
|
|
|
|
|
# ── build_models: conditioning.share_stages ─────────────────────────────────
|
|
|
|
|
|
def _minimal_model_config(share_stages: bool) -> dict:
|
|
cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
|
|
cfg["conditioning"]["share_stages"] = share_stages
|
|
cfg["conditioning"]["particle"]["emb_dim"] = 4
|
|
cfg["conditioning"]["material"]["emb_dim"] = 4
|
|
cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
|
|
cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
|
|
return {
|
|
"pdg_vocab": 3,
|
|
"mat_vocab": 2,
|
|
"conditioning": cfg["conditioning"],
|
|
"stage1_model": cfg["stage1_model"],
|
|
"stage2_model": cfg["stage2_model"],
|
|
}
|
|
|
|
|
|
def test_build_models_share_stages_true_shares_condition_encoder_instance():
|
|
built = build_models(_minimal_model_config(share_stages=True))
|
|
stage1, stage2 = built["stage1"], built["stage2"]
|
|
assert stage1 is not None and stage2 is not None
|
|
assert stage1.cond_enc is stage2.cond_enc
|
|
|
|
|
|
def test_build_models_share_stages_false_builds_independent_condition_encoders():
|
|
built = build_models(_minimal_model_config(share_stages=False))
|
|
stage1, stage2 = built["stage1"], built["stage2"]
|
|
assert stage1 is not None and stage2 is not None
|
|
assert stage1.cond_enc is not stage2.cond_enc
|
|
|
|
|
|
def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
|
|
"""The shared encoder's parameters must actually appear in both stages'
|
|
own `.parameters()` — that's what makes each stage's independent
|
|
optimizer include (and update) them, which is the actual mechanism behind
|
|
"shared weights, forced common representation", not just object identity
|
|
on `.cond_enc`."""
|
|
built = build_models(_minimal_model_config(share_stages=True))
|
|
stage1, stage2 = built["stage1"], built["stage2"]
|
|
assert stage1 is not None and stage2 is not None
|
|
|
|
shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
|
|
assert shared_ids
|
|
assert shared_ids <= {id(p) for p in stage1.parameters()}
|
|
assert shared_ids <= {id(p) for p in stage2.parameters()}
|
|
|
|
|
|
def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
|
|
"""gitea #29 end-to-end through build_models: setting
|
|
stage2_model.particle_type.n_classes independently of
|
|
conditioning.particle.emb_dim actually resizes the built stage2 model,
|
|
not just the two lower-level unit tests above."""
|
|
cfg = _minimal_model_config(share_stages=False) # conditioning.particle.emb_dim = 4
|
|
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
|
|
built = build_models(cfg)
|
|
assert built["stage2"] is not None
|
|
assert built["stage2"].type_dim == 11
|
|
|
|
|
|
def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
|
|
cfg = _minimal_model_config(share_stages=False) # conditioning.particle.emb_dim = 4
|
|
cfg["stage2_model"]["generator"] = "wgan"
|
|
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
|
|
default_n_classes_critic = build_critics(cfg)["stage2"]
|
|
assert default_n_classes_critic is not None
|
|
|
|
cfg["stage2_model"]["particle_type"]["n_classes"] = 11
|
|
wider_critic = build_critics(cfg)["stage2"]
|
|
assert wider_critic is not None
|
|
# k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
|
|
# widening n_classes alone (emb_dim stays 4) must widen the critic input.
|
|
assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
|
|
|
|
|
|
# ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
|
|
|
|
|
|
def _partial_model_config() -> dict:
|
|
"""A hand-built model_config that omits stage2_model.decoder and
|
|
stage2_model.particle_type — deliberately not derived from
|
|
DEFAULT_CONFIG, unlike _minimal_model_config above. Regression fixture
|
|
for issues.md Issue 1: build_models/build_critics/StageSpec.from_config's
|
|
own fallback defaults for these two keys must equal DEFAULT_CONFIG's
|
|
("autoregressive" / "onehot"), not the old, now-wrong v0.2-shaped
|
|
("one_shot" / "physical") literals that used to live in three separate
|
|
.get(key, default) call sites."""
|
|
return {
|
|
"pdg_vocab": 3,
|
|
"mat_vocab": 2,
|
|
"conditioning": {
|
|
"particle": {"type": "physical", "emb_dim": 4, "n_layers": 1},
|
|
"material": {"type": "physical", "emb_dim": 4, "n_layers": 1},
|
|
},
|
|
"stage1_model": {"active": False},
|
|
"stage2_model": {
|
|
"generator": "wgan",
|
|
"hidden_dim": 8,
|
|
"n_res_blocks": 1,
|
|
"k_max": 3,
|
|
# decoder and particle_type deliberately omitted
|
|
},
|
|
}
|
|
|
|
|
|
def test_build_models_omitted_decoder_and_particle_type_match_default_config():
|
|
built = build_models(_partial_model_config())
|
|
assert isinstance(built["stage2"], Stage2Autoregressive)
|
|
assert built["stage2"].particle_type_cfg["target"] == "onehot"
|
|
|
|
|
|
def test_build_models_custom_heads_block_controls_head_shapes():
|
|
"""gitea #36: stage{1,2}_model.heads flows all the way from config dict
|
|
through build_models to the actual constructed head shapes."""
|
|
cfg = _partial_model_config()
|
|
cfg["stage1_model"] = {
|
|
"active": True,
|
|
"hidden_dim": 40,
|
|
"n_res_blocks": 1,
|
|
"heads": {"n_sec": {"hidden_ratio": 0.25, "depth": 1}},
|
|
}
|
|
cfg["stage2_model"]["decoder"] = "one_shot"
|
|
cfg["stage2_model"]["generator"] = "flow" # wgan folds the type slice; no separate type_head
|
|
cfg["stage2_model"]["n_sec"] = {"owner": "stage1"}
|
|
cfg["stage2_model"]["heads"] = {
|
|
"n_sec": {"hidden_ratio": 0.25, "depth": 1},
|
|
"type": {"hidden_ratio": 0.75, "depth": 2},
|
|
}
|
|
|
|
built = build_models(cfg)
|
|
stage1, stage2 = built["stage1"], built["stage2"]
|
|
assert stage1 is not None
|
|
assert stage2 is not None
|
|
|
|
assert stage1.n_sec_head is not None
|
|
assert len(stage1.n_sec_head) == 1 # owner=stage1, so stage1 builds it
|
|
assert stage2.n_sec_head is None # owner=stage1, so stage2 doesn't
|
|
|
|
assert isinstance(stage2, Stage2OneShot)
|
|
assert stage2.type_head is not None
|
|
assert len(stage2.type_head) == 3
|
|
assert stage2.type_head[0].out_features == 6 # round(8 * 0.75)
|
|
|
|
|
|
def test_build_critics_omitted_particle_type_matches_default_config():
|
|
cfg = _partial_model_config()
|
|
cfg["stage2_model"]["generator"] = "wgan"
|
|
onehot_critic = build_critics(cfg)["stage2"]
|
|
assert onehot_critic is not None
|
|
onehot_in_dim = onehot_critic.input_proj.in_features
|
|
|
|
cfg["stage2_model"]["particle_type"] = {"target": "physical"}
|
|
physical_critic = build_critics(cfg)["stage2"]
|
|
assert physical_critic is not None
|
|
physical_in_dim = physical_critic.input_proj.in_features
|
|
|
|
# onehot's per-slot type width is emb_dim classes vs. physical's fixed
|
|
# (log-mass, charge) pair — different unless emb_dim happens to be 2, so
|
|
# this also confirms the critic was actually built in onehot mode by
|
|
# default, not silently falling back to physical.
|
|
assert onehot_in_dim != physical_in_dim
|
|
|
|
|
|
# ── build_critics: critic_hidden_dim/critic_n_res_blocks honoured (gitea #28) ─
|
|
|
|
|
|
def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_generator_size():
|
|
cfg = _minimal_model_config(share_stages=False)
|
|
cfg["stage1_model"]["generator"] = "wgan"
|
|
cfg["stage1_model"]["hidden_dim"] = 8
|
|
cfg["stage1_model"]["n_res_blocks"] = 1
|
|
|
|
inherited = build_critics(cfg)["stage1"]
|
|
assert inherited is not None
|
|
assert inherited.input_proj.out_features == 8
|
|
assert len(inherited.blocks) == 1
|
|
|
|
cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
|
|
cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
|
|
overridden = build_critics(cfg)["stage1"]
|
|
assert overridden is not None
|
|
assert overridden.input_proj.out_features == 16
|
|
assert len(overridden.blocks) == 3
|
|
|
|
|
|
def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
|
|
cfg = _minimal_model_config(share_stages=False)
|
|
cfg["stage2_model"]["generator"] = "wgan"
|
|
cfg["stage2_model"]["hidden_dim"] = 8
|
|
cfg["stage2_model"]["n_res_blocks"] = 1
|
|
|
|
inherited = build_critics(cfg)["stage2"]
|
|
assert inherited is not None
|
|
assert inherited.input_proj.out_features == 8
|
|
assert len(inherited.blocks) == 1
|
|
|
|
cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
|
|
cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
|
|
overridden = build_critics(cfg)["stage2"]
|
|
assert overridden is not None
|
|
assert overridden.input_proj.out_features == 16
|
|
assert len(overridden.blocks) == 3
|