lars
da7cde3ef9
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 36s
CI / Type check (ty) (push) Successful in 39s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 30s
CI / Tests (pull_request) Successful in 2m50s
CI / Tests (push) Successful in 2m58s
Works through docs/v0.3.0-followups.md item by item, closing the gap between the design doc and the shipped v0.3.0-stage2-autoregressive code: 1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2 dispatch, stage-2 particle-type-class marginal. 2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run. 3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/ train instead of the hardcoded K_MAX constant. 4. Mixed conditioning.particle.type / conditioning.material.type support end-to-end (data pipeline + dwarf warm-cache). 5. conditioning.share_stages = true: one shared ConditionEncoder instance across both stages. 6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2 (was silently unimplemented). 7. giant predict/rollout: implement conditioning.*.type = "onehot" via the checkpoint's saved pdg_topn_map/mat_topn_map. 8. network.py's checkpoint-path model_config migration now fails loudly on non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's TOML-load path (§4.2). 9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a rollout-capable checkpoint (§9). Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics), mostly a test-helper dict-unpack pattern that made every unrelated constructor keyword look like a type error. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
import torch
|
|
from giant.constants import COND_DIM
|
|
from giant.model.network import Stage1Model
|
|
from giant.model.schedule import CosineSchedule, flow_matching_loss
|
|
from giant.sample import sample_flow, sample_ddim
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
|
|
|
|
def _small_model():
|
|
return 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,
|
|
)
|
|
|
|
|
|
def _batch(B=8):
|
|
x1 = torch.randn(B, 9)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
return x1, cond_cont, cond_cat
|
|
|
|
|
|
def test_flow_matching_loss_nonneg():
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_flow_matching_loss_is_scalar():
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
|
|
|
|
def test_flow_matching_loss_has_grad():
|
|
model = _small_model()
|
|
x1, cond_cont, cond_cat = _batch()
|
|
flow_matching_loss(model, x1, cond_cont, cond_cat).backward()
|
|
assert any(p.grad is not None for p in model.parameters())
|
|
|
|
|
|
def test_sample_flow_shape():
|
|
B = 6
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
|
assert sample.shape == (B, 9)
|
|
assert n_sec is not None and n_sec.shape == (B,)
|
|
|
|
|
|
def test_ddpm_loss_nonneg():
|
|
schedule = CosineSchedule(T=50)
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = schedule.loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_sample_ddim_shape():
|
|
B = 4
|
|
schedule = CosineSchedule(T=50)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
|
assert sample.shape == (B, 9)
|
|
assert n_sec is not None and n_sec.shape == (B,)
|