lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 1 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_sample.py
+10 -30
View File
@@ -30,9 +30,7 @@ def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]
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
)
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
return cond_cont, cond_cat
@@ -43,12 +41,8 @@ def _conditioning_for(target: str) -> str:
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
)
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 = {"target": target}
# build_models (giant/model/network.py) computes sec_dim this same way
# before constructing Stage2OneShot — its own default (SEC_DIM, the
@@ -78,9 +72,7 @@ def _stage2_ar(
k_max: int = 5,
history: str = "markov",
) -> Stage2Autoregressive:
particle_cfg, material_cfg = _particle_material_cfg(
_conditioning_for(target), emb_dim
)
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
return Stage2Autoregressive(
pdg_vocab=pdg,
mat_vocab=mat,
@@ -163,9 +155,7 @@ def test_sample_secondaries_flow_shapes_by_target(target):
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
)
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)
@@ -180,9 +170,7 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
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
)
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)
@@ -196,15 +184,11 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
@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
)
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
)
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)
@@ -220,9 +204,7 @@ def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
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
)
_, _, 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()
@@ -237,8 +219,6 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
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
)
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]]