lars
e6e0eb22bf
Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs); Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1 output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
import torch
|
|
from giant.constants import COND_DIM
|
|
from giant.model.network import DenoisingMLP, SinusoidalEmbedding
|
|
|
|
|
|
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_denoising_mlp_output_shape():
|
|
B = 8
|
|
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
|
|
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, t, cond_cont, cond_cat)
|
|
assert out.shape == (B, 9)
|
|
|
|
|
|
def test_denoising_mlp_gradients_flow():
|
|
B = 4
|
|
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
|
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, t, cond_cont, cond_cat).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}"
|