v0.3.0 step 2: network.py refactor to composable stage models

Decomposes the ten permutation classes in giant/model/network.py into
the reusable parts from docs/v0.3.0-design.md §5: ConditionEncoder (now
independently configurable per particle/material axis), ContextAdapter,
Trunk/MonolithicTrunk/RoutedTrunk/ExpertTrunk, and the stage classes
Stage1Model/Stage2OneShot/CriticModel (Stage2Autoregressive stubbed,
raises NotImplementedError until step 4/5). build_models/build_critics
now return a dict keyed by stage and accept the new nested config shape,
with routed WGAN reachable for the first time (the old --mode wgan
--router rejection is gone) and stage2_model.router.tie_to_stage1
sharing a literal Router instance.

A v0.2 checkpoint's flat model_config auto-migrates via
_migrate_legacy_model_config + migrate_legacy_state_dict, preserving the
n_sec_head's attachment to Stage1Model (legacy_owner="stage1", design
doc §4.1). tests/test_migration_v02_v03.py proves this bit-identical
against a frozen v0.2 snapshot (tests/legacy/network_v02_snapshot.py)
for both flow and wgan, both conditioning modes.
scripts/check_migration_v02_v03.py is the real-checkpoint counterpart
for a portal machine with /ceph access.

giant/model/schedule.py's flow-matching/DDPM loss helpers are updated
to the new model-call convention (t as a keyword). giant/sample.py,
giant/rollout.py, and giant/validate.py are not yet updated (deferred
to design doc step 6) — their exercising tests are marked xfail with
that reasoning rather than silently broken.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_phase2.py

@@ -5,31 +5,50 @@ import pytest
 import torch
 
 from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM
-from giant.model.network import DenoisingMLP, SecondaryDecoder
+from giant.model.network import Stage1Model, Stage2OneShot
 from giant.model.schedule import flow_matching_loss_secondary
 from giant.sample import sample_secondaries
 
+_SAMPLE_SECONDARIES_XFAIL_REASON = (
+    "giant/sample.py isn't updated yet — sample_secondaries calls the "
+    "decoder positionally as decoder(x, t, cond_cont, cond_cat, stage1_out), "
+    "which doesn't match Stage2OneShot's new forward signature. Deferred to "
+    "docs/v0.3.0-design.md step 6."
+)
+
 
 # ── helpers ──────────────────────────────────────────────────────────────────
 
 
+def _particle_material_cfg(conditioning: str) -> tuple[dict, dict]:
+    cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
+    return dict(cfg), dict(cfg)
+
+
 def _stage1(pdg=3, mat=2, conditioning="embedding"):
-    return DenoisingMLP(
+    particle_cfg, material_cfg = _particle_material_cfg(conditioning)
+    return Stage1Model(
         pdg_vocab=pdg,
         mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
         hidden_dim=32,
-        n_blocks=2,
-        conditioning=conditioning,
+        n_res_blocks=2,
+        n_sec_head_k_max=K_MAX,
     )
 
 
 def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
-    return SecondaryDecoder(
+    particle_cfg, material_cfg = _particle_material_cfg(conditioning)
+    return Stage2OneShot(
         pdg_vocab=pdg,
         mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
         hidden_dim=32,
-        n_blocks=2,
-        conditioning=conditioning,
+        n_res_blocks=2,
+        generator="flow",
+        time_dim=16,
     )
 
 
@@ -41,7 +60,7 @@ def _cond(B=8, pdg=3, mat=2):
     return cond_cont, cond_cat
 
 
-# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
+# ── Stage1Model Phase-2 additions ────────────────────────────────────────────
 
 
 def test_predict_n_sec_shape():
@@ -82,7 +101,7 @@ def test_condition_encoder_embedding_mode_has_embedding_tables():
     assert not hasattr(model.cond_enc, "particle_mlp")
 
 
-# ── SecondaryDecoder ──────────────────────────────────────────────────────────
+# ── Stage2OneShot ─────────────────────────────────────────────────────────────
 
 
 @pytest.mark.parametrize("conditioning", ["embedding", "physical"])
@@ -93,7 +112,7 @@ def test_sec_decoder_output_shape(conditioning):
     t = torch.rand(B)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
-    out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
+    out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
     assert out.shape == (B, SEC_DIM)
 
 
@@ -104,7 +123,7 @@ def test_sec_decoder_no_nan():
     t = torch.rand(B)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
-    out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
+    out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
     assert torch.isfinite(out).all()
 
 
@@ -115,7 +134,9 @@ def test_sec_decoder_gradients():
     t = torch.rand(B)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
-    decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
+    flow_out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t).sum()
+    nsec_out = decoder.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
+    (flow_out + nsec_out).backward()
     for name, p in decoder.named_parameters():
         assert p.grad is not None, f"no grad for {name}"
 
@@ -167,6 +188,7 @@ def test_flow_matching_loss_secondary_has_grad():
 # ── sampling ──────────────────────────────────────────────────────────────────
 
 
+@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False)
 def test_sample_secondaries_shapes():
     B, pdg, mat = 6, 3, 2
     decoder = _sec_decoder(pdg, mat)
@@ -182,6 +204,7 @@ def test_sample_secondaries_shapes():
     assert sec_valid.dtype == torch.bool
 
 
+@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False)
 def test_sample_secondaries_valid_mask_matches_n_sec():
     B, pdg, mat = 4, 3, 2
     decoder = _sec_decoder(pdg, mat)