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

@@ -4,9 +4,19 @@ import numpy as np
 import pytest
 import torch
 
-from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM
-from giant.model.network import Stage1Model, Stage2OneShot
-from giant.model.schedule import flow_matching_loss_secondary
+from giant.constants import (
+    COND_DIM,
+    CONT_SLOT_DIM,
+    K_MAX,
+    PARTICLE_PHYS_DIM,
+    SEC_DIM,
+    X_DIM,
+)
+from giant.model.network import Stage1Model, Stage2Autoregressive, Stage2OneShot
+from giant.model.schedule import (
+    flow_matching_loss_secondary,
+    flow_matching_loss_secondary_ar,
+)
 from giant.sample import sample_secondaries
 
 _SAMPLE_SECONDARIES_XFAIL_REASON = (
@@ -185,6 +195,108 @@ def test_flow_matching_loss_secondary_has_grad():
     assert any(p.grad is not None for p in decoder.parameters())
 
 
+# ── masked flow matching loss — autoregressive (v0.3.0 step 5) ─────────────
+
+
+def _sec_decoder_ar(pdg=3, mat=2, k_max=K_MAX):
+    particle_cfg, material_cfg = _particle_material_cfg("embedding")
+    return Stage2Autoregressive(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator="flow",
+        time_dim=16,
+        k_max=k_max,
+    )
+
+
+def _ar_history_inputs(B, K, hist_dim):
+    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_flow_matching_loss_secondary_ar_scalar():
+    B, K, pdg, mat = 8, K_MAX, 3, 2
+    decoder = _sec_decoder_ar(pdg, mat, k_max=K)
+    x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
+    cond_cont, cond_cat = _cond(B, pdg, mat)
+    stage1_out = torch.randn(B, X_DIM)
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
+        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    )
+    sec_mask = torch.ones(B, K, dtype=torch.bool)
+    loss = flow_matching_loss_secondary_ar(
+        decoder,
+        x1,
+        cond_cont,
+        cond_cat,
+        stage1_out,
+        history_feat,
+        has_prev,
+        remaining_frac,
+        slot_idx,
+        sec_mask,
+    )
+    assert loss.shape == ()
+    assert loss.item() >= 0.0
+
+
+def test_flow_matching_loss_secondary_ar_mask_zeros_padding():
+    B, K, pdg, mat = 4, K_MAX, 3, 2
+    decoder = _sec_decoder_ar(pdg, mat, k_max=K)
+    x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
+    cond_cont, cond_cat = _cond(B, pdg, mat)
+    stage1_out = torch.randn(B, X_DIM)
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
+        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    )
+    sec_mask = torch.zeros(B, K, dtype=torch.bool)
+    loss = flow_matching_loss_secondary_ar(
+        decoder,
+        x1,
+        cond_cont,
+        cond_cat,
+        stage1_out,
+        history_feat,
+        has_prev,
+        remaining_frac,
+        slot_idx,
+        sec_mask,
+    )
+    assert loss.item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_flow_matching_loss_secondary_ar_has_grad():
+    B, K, pdg, mat = 4, K_MAX, 3, 2
+    decoder = _sec_decoder_ar(pdg, mat, k_max=K)
+    x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
+    cond_cont, cond_cat = _cond(B, pdg, mat)
+    stage1_out = torch.randn(B, X_DIM)
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
+        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    )
+    sec_mask = torch.ones(B, K, dtype=torch.bool)
+    flow_matching_loss_secondary_ar(
+        decoder,
+        x1,
+        cond_cont,
+        cond_cat,
+        stage1_out,
+        history_feat,
+        has_prev,
+        remaining_frac,
+        slot_idx,
+        sec_mask,
+    ).backward()
+    assert any(p.grad is not None for p in decoder.parameters())
+
+
 # ── sampling ──────────────────────────────────────────────────────────────────