v0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation

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>

tests/test_network.py

@@ -1,9 +1,12 @@
+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,
@@ -290,3 +293,228 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
     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}"