lars
c9d255b1c5
CI / Format (ruff format) (push) Successful in 30s
CI / Lint (ruff check) (push) Successful in 31s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 33s
CI / Lint (ruff check) (pull_request) Successful in 33s
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 31s
CI / Tests (push) Successful in 2m12s
CI / Tests (pull_request) Successful in 2m10s
Replaces the Stage2Autoregressive stub with a real per-token secondary decoder: MarkovHistory summarizes the previous secondary, remaining-energy fraction and slot index round out the per-token conditioning, and the existing Trunk/MonolithicTrunk/RoutedTrunk machinery is reused unchanged by batching all K_MAX tokens together under teacher forcing (one parallel pass, no new trunk code). build_models/build_critics wire it in; the WGAN critic stays whole-sequence, so build_critics needs no AR-specific path. train.py's FlowDDPMStageTrainer/WGANStageTrainer gain a decoder branch, sharing optimizer/EMA/checkpoint machinery with the one-shot path. _assemble_stage2_real is now defined in terms of the new unflattened _assemble_stage2_ar_target helper, removing a near-duplicate branch. Also lands the §11.4 differentiability validation-obligation instrumentation (trunk-gradient norm from the particle-type slice vs. the continuous slices, for generator=wgan + particle_type.target=onehot) via backward hooks in _relax_onehot_type_slice, decoder-agnostic and surfaced as two new metrics.csv columns. This also fixes the standing regression where any config not explicitly overriding decoder="one_shot" crashed at build_models, since stage2_model.decoder defaults to "autoregressive" — confirmed by removing tests/test_pipeline.py's now-stale override so the default config runs end-to-end against real synthetic data. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
521 lines
16 KiB
Python
521 lines
16 KiB
Python
import pytest
|
|
import torch
|
|
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
|
|
from giant.model.network import (
|
|
ConditionEncoder,
|
|
MarkovHistory,
|
|
SinusoidalEmbedding,
|
|
Stage1Model,
|
|
Stage2Autoregressive,
|
|
Stage2OneShot,
|
|
cat_col_layout,
|
|
stage2_trunk_sec_dim,
|
|
stage2_type_dim,
|
|
)
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
|
|
ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
|
|
|
|
|
|
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_stage1_model_output_shape():
|
|
B = 8
|
|
model = Stage1Model(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
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, cond_cont, cond_cat, t=t)
|
|
assert out.shape == (B, 9)
|
|
|
|
|
|
def test_stage1_model_gradients_flow():
|
|
B = 4
|
|
model = 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,
|
|
)
|
|
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, cond_cont, cond_cat, t=t).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}"
|
|
|
|
|
|
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 —
|
|
decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2."""
|
|
model = Stage1Model(
|
|
pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
|
|
)
|
|
assert model.n_sec_head is None
|
|
|
|
|
|
# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
|
|
|
|
|
|
def test_cat_col_layout_neither_onehot():
|
|
assert cat_col_layout("physical", "embedding") == (None, None)
|
|
|
|
|
|
def test_cat_col_layout_particle_only():
|
|
assert cat_col_layout("onehot", "physical") == (2, None)
|
|
|
|
|
|
def test_cat_col_layout_material_only():
|
|
assert cat_col_layout("physical", "onehot") == (None, 2)
|
|
|
|
|
|
def test_cat_col_layout_both_onehot_particle_then_material():
|
|
assert cat_col_layout("onehot", "onehot") == (2, 3)
|
|
|
|
|
|
def test_stage2_type_dim_physical_is_particle_phys_dim():
|
|
assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
|
|
|
|
|
|
def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
|
|
assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
|
|
assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
|
|
|
|
|
|
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
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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
|
|
)
|
|
|
|
|
|
# --- ConditionEncoder onehot mode -------------------------------------------
|
|
|
|
|
|
def test_condition_encoder_onehot_forward_shape_and_gradients():
|
|
B = 8
|
|
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
|
material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
|
|
enc = ConditionEncoder(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=ONEHOT_PARTICLE_CFG,
|
|
material_cfg=ONEHOT_MATERIAL_CFG,
|
|
out_dim=32,
|
|
)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.randint(0, 5, (B,)),
|
|
torch.randint(0, 3, (B,)),
|
|
torch.randint(0, particle_emb_dim, (B,)),
|
|
torch.randint(0, material_emb_dim, (B,)),
|
|
],
|
|
dim=1,
|
|
)
|
|
out = enc(cond_cont, cond_cat)
|
|
assert out.shape == (B, 32)
|
|
# onehot itself is unlearned, but the fusion MLP downstream still has
|
|
# gradients — the encoder as a whole must still be trainable.
|
|
out.sum().backward()
|
|
assert enc.mlp[0].weight.grad is not None
|
|
|
|
|
|
def test_condition_encoder_onehot_is_a_true_one_hot_vector():
|
|
"""The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
|
|
verify the concatenated input segment really is one-hot, not e.g. an
|
|
accidentally-learned embedding."""
|
|
B = 4
|
|
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
|
enc = ConditionEncoder(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=ONEHOT_PARTICLE_CFG,
|
|
material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
|
|
out_dim=16,
|
|
)
|
|
cond_cont = torch.zeros(B, COND_DIM)
|
|
idx = torch.tensor([0, 1, 2, 5])
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.zeros(B, dtype=torch.long),
|
|
torch.zeros(B, dtype=torch.long),
|
|
idx.clamp(max=particle_emb_dim - 1),
|
|
],
|
|
dim=1,
|
|
)
|
|
pdg_e = enc._particle_embed(cond_cont, cond_cat)
|
|
assert pdg_e.shape == (B, particle_emb_dim)
|
|
assert torch.all(pdg_e.sum(dim=-1) == 1.0)
|
|
|
|
|
|
# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) --
|
|
|
|
|
|
def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
|
|
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
|
if target != "physical":
|
|
particle_cfg = dict(particle_cfg)
|
|
if target == "embedding":
|
|
particle_cfg["type"] = "embedding"
|
|
k_max = 5
|
|
sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
|
|
return Stage2OneShot(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=16,
|
|
n_res_blocks=1,
|
|
cond_out_dim=16,
|
|
context_dim=8,
|
|
sec_dim=sec_dim,
|
|
generator=generator,
|
|
k_max=k_max,
|
|
particle_type_cfg={"target": target, "lambda": 1.0},
|
|
)
|
|
|
|
|
|
def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
|
|
assert _build_stage2("physical", "flow").type_head is None
|
|
assert _build_stage2("physical", "wgan").type_head is None
|
|
|
|
|
|
def test_stage2_oneshot_onehot_flow_has_type_head():
|
|
model = _build_stage2("onehot", "flow")
|
|
assert model.type_head is not None
|
|
|
|
|
|
def test_stage2_oneshot_onehot_wgan_has_no_type_head():
|
|
"""Under wgan the type slice is folded into forward()'s own output and
|
|
relaxed via ST-Gumbel by the trainer — no separate head needed."""
|
|
model = _build_stage2("onehot", "wgan")
|
|
assert model.type_head is None
|
|
|
|
|
|
def test_stage2_oneshot_embedding_flow_has_type_head():
|
|
model = _build_stage2("embedding", "flow")
|
|
assert model.type_head is not None
|
|
|
|
|
|
def test_stage2_oneshot_predict_type_shape():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
out = model.predict_type(cond_cont, cond_cat, stage1_out)
|
|
assert out.shape == (B, k_max, emb_dim)
|
|
|
|
|
|
def test_stage2_oneshot_predict_type_raises_when_no_type_head():
|
|
model = _build_stage2("physical", "flow")
|
|
cond_cont = torch.randn(2, COND_DIM)
|
|
cond_cat = torch.zeros(2, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(2, 9)
|
|
try:
|
|
model.predict_type(cond_cont, cond_cat, stage1_out)
|
|
raise AssertionError("expected RuntimeError")
|
|
except RuntimeError:
|
|
pass
|
|
|
|
|
|
def test_stage2_oneshot_forward_shape_onehot_wgan():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
z = torch.randn(B, model.noise_dim)
|
|
out = model(z, cond_cont, cond_cat, stage1_out)
|
|
assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))
|
|
|
|
|
|
def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
|
|
t = torch.rand(B)
|
|
out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
|
assert out.shape == (B, k_max * CONT_SLOT_DIM)
|
|
|
|
|
|
# --- MarkovHistory (docs/v0.3.0-design.md §6.2) -----------------------------
|
|
|
|
|
|
def test_markov_history_shape():
|
|
hist = MarkovHistory(in_dim=7, out_dim=12)
|
|
B, K = 3, 5
|
|
feat = torch.randn(B, K, 7)
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
out = hist(feat, has_prev)
|
|
assert out.shape == (B, K, 12)
|
|
|
|
|
|
def test_markov_history_uses_start_vector_when_no_prev():
|
|
"""Slot 0's own raw feature must be ignored — a learned start vector is
|
|
substituted there instead (a reasonable default not specified by the
|
|
design doc, see Stage2Autoregressive's docstring)."""
|
|
hist = MarkovHistory(in_dim=4, out_dim=6)
|
|
B, K = 2, 3
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
feat_a = torch.randn(B, K, 4)
|
|
feat_b = feat_a.clone()
|
|
feat_b[:, 0] = torch.randn(B, 4) * 100
|
|
out_a = hist(feat_a, has_prev)
|
|
out_b = hist(feat_b, has_prev)
|
|
assert torch.allclose(out_a[:, 0], out_b[:, 0])
|
|
assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
|
|
|
|
|
|
# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) ---------
|
|
|
|
|
|
def _build_stage2_ar(
|
|
target: str,
|
|
generator: str,
|
|
emb_dim: int = 6,
|
|
k_max: int = 5,
|
|
history: str = "markov",
|
|
) -> Stage2Autoregressive:
|
|
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
|
if target == "embedding":
|
|
particle_cfg = dict(particle_cfg)
|
|
particle_cfg["type"] = "embedding"
|
|
return Stage2Autoregressive(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=16,
|
|
n_res_blocks=1,
|
|
cond_out_dim=16,
|
|
context_dim=8,
|
|
generator=generator,
|
|
k_max=k_max,
|
|
particle_type_cfg={"target": target, "lambda": 1.0},
|
|
history=history,
|
|
)
|
|
|
|
|
|
def _ar_inputs(B: int, K: int, hist_dim: int):
|
|
history_feat = torch.randn(B, K, hist_dim)
|
|
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
|
remaining_frac = torch.rand(B, K)
|
|
slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1)
|
|
return history_feat, has_prev, remaining_frac, slot_idx
|
|
|
|
|
|
def test_stage2_autoregressive_history_attention_raises():
|
|
with pytest.raises(NotImplementedError):
|
|
_build_stage2_ar("onehot", "wgan", history="attention")
|
|
|
|
|
|
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
|
@pytest.mark.parametrize("generator", ["wgan", "flow"])
|
|
def test_stage2_autoregressive_forward_shape(target, generator):
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
|
|
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
|
|
)
|
|
token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
|
|
if generator == "wgan":
|
|
x_t = torch.randn(B, K, model.noise_dim)
|
|
t = None
|
|
else:
|
|
x_t = torch.randn(B, K, token_dim)
|
|
t = torch.rand(B, K)
|
|
out = model(
|
|
x_t,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
t=t,
|
|
)
|
|
assert out.shape == (B, K, token_dim)
|
|
|
|
|
|
def test_stage2_autoregressive_predict_n_sec_shape():
|
|
B, k_max, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=k_max)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
stage1_out = torch.randn(B, 9)
|
|
logits = model.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
|
assert logits.shape == (B, k_max + 1)
|
|
|
|
|
|
def test_stage2_autoregressive_predict_type_shape():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
|
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": "onehot"}, emb_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,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
)
|
|
assert out.shape == (B, K, emb_dim)
|
|
|
|
|
|
@pytest.mark.parametrize("target,generator", [("physical", "flow"), ("onehot", "wgan")])
|
|
def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, generator):
|
|
B, K, emb_dim = 2, 5, 6
|
|
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
|
|
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
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
model.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
)
|
|
|
|
|
|
def test_stage2_autoregressive_gradients_flow_wgan_onehot():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "wgan", emb_dim=emb_dim, k_max=K)
|
|
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": "onehot"}, emb_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,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
).sum()
|
|
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
|
(gen_out + nsec_out).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
def test_stage2_autoregressive_gradients_flow_onehot():
|
|
B, K, emb_dim = 4, 5, 6
|
|
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
|
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": "onehot"}, emb_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)
|
|
flow_out = model(
|
|
x_t,
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
t=t,
|
|
).sum()
|
|
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
|
type_out = model.predict_type(
|
|
cond_cont,
|
|
cond_cat,
|
|
stage1_out,
|
|
history_feat,
|
|
has_prev,
|
|
remaining_frac,
|
|
slot_idx,
|
|
).sum()
|
|
(flow_out + nsec_out + type_out).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|