lars
c984d0a19d
CI / Sync project version with tag (hand-pushed tags only) (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 53s
CI / Type check (ty) (pull_request) Successful in 57s
CI / Lint (ruff check) (pull_request) Successful in 41s
CI / Tests (pull_request) Successful in 8m20s
CI / Release (bump, changelog, badges, tag) on merge to master (pull_request) Has been skipped
uv.lock was stale (ty 0.0.50 -> 0.0.78, ruff 0.15 -> 0.16, polars, numpy, typer, wandb, pytest, and others), all within existing pyproject.toml bounds. ruff 0.16 widened its default rule selection, taking this repo from 0 to 274 lint errors under the same config; --fix handled most of it (import sorting, Optional[X] -> X | None, ...), and the remainder (unused unpacked variables, dict()-as-literal, subprocess.run without explicit check=, a couple of intentional broad excepts/naive datetimes) were fixed or annotated by hand. Also fixes a real type-narrowing gap ty 0.0.78 caught in test_config_consumed_keys.py's `or`-combined isinstance check. torch stays pinned to 2.3.x (deliberate, see CLAUDE.md); pyarrow's <25 ceiling is left as a separate decision. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TMdZFqXXig7i3XkirSUxef
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import torch
|
|
|
|
from giant.config import ConditioningAxisConfig
|
|
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_ddim, sample_flow
|
|
|
|
PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
|
|
MATERIAL_CFG = ConditioningAxisConfig(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,)
|