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_network.py

@@ -1,6 +1,9 @@
 import torch
 from giant.constants import COND_DIM
-from giant.model.network import DenoisingMLP, SinusoidalEmbedding
+from giant.model.network import SinusoidalEmbedding, Stage1Model
+
+PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 
 
 def test_sinusoidal_embedding_shape():
@@ -15,9 +18,15 @@ def test_sinusoidal_embedding_batch_1():
     assert emb(t).shape == (1, 32)
 
 
-def test_denoising_mlp_output_shape():
+def test_stage1_model_output_shape():
     B = 8
-    model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
+    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)
@@ -28,20 +37,37 @@ def test_denoising_mlp_output_shape():
         ],
         dim=1,
     )
-    out = model(x_t, t, cond_cont, cond_cat)
+    out = model(x_t, cond_cont, cond_cat, t=t)
     assert out.shape == (B, 9)
 
 
-def test_denoising_mlp_gradients_flow():
+def test_stage1_model_gradients_flow():
     B = 4
-    model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
+    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, t, cond_cont, cond_cat).sum()
+    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