lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Bump ruff line-length to 120 and reformat
CI / Lint (ruff check) (push) Successful in 31s
Details
CI / Format (ruff format) (push) Successful in 32s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Lint (ruff check) (pull_request) Successful in 26s
Details
CI / Type check (ty) (push) Successful in 29s
Details
CI / Format (ruff format) (pull_request) Successful in 33s
Details
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 35s
Details
CI / Tests (pull_request) Successful in 3m47s
Details
CI / Tests (push) Successful in 3m55s
Details

Browse Source 
Rejoins lines that only wrapped because they exceeded the old 88-char
limit; ruff check and the full test suite (725 passed) are unaffected.
This commit is contained in:
Lars Bogner
2026-08-12 13:33:09 +02:00
parent 9ce7b32324
commit 55332db67a
66 changed files with 757 additions and 2413 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_network.py
+13 -42
View File
@@ -85,9 +85,7 @@ def test_stage1_model_gradients_flow():
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
)
model = Stage1Model(pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG)
assert model.n_sec_head is None
@@ -121,29 +119,18 @@ def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
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
)
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)
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
)
assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
# --- ConditionEncoder onehot mode -------------------------------------------
@@ -440,16 +427,12 @@ def test_stage2_autoregressive_history_invalid_raises():
@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
)
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
)
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)
@@ -488,9 +471,7 @@ def test_stage2_autoregressive_predict_type_shape():
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
)
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,
@@ -511,9 +492,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
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
)
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,
@@ -533,9 +512,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
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
)
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,
@@ -560,9 +537,7 @@ def test_stage2_autoregressive_gradients_flow_onehot():
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
)
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)
@@ -601,17 +576,13 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
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 = _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
)
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)