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

@@ -1,18 +1,36 @@
 """Tests for giant/train.py."""
 
 import copy
+import csv
 import tempfile
 from pathlib import Path
 
 import pytest
 import torch
 
-from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
+from giant.constants import (
+    COND_DIM,
+    CONT_SLOT_DIM,
+    K_MAX,
+    PARTICLE_PHYS_DIM,
+    SEC_SLOT_DIM,
+    X_DIM,
+)
 from giant.model.network import build_critics, build_models
 from giant.train import (
     FlowDDPMStageTrainer,
     WGANStageTrainer,
+    _ar_has_prev,
+    _assemble_stage2_ar_inputs,
+    _assemble_stage2_ar_target,
+    _assemble_stage2_real,
+    _build_stage_trainers,
     _gumbel_tau,
+    _relax_onehot_type_slice,
+    _remaining_energy_fraction,
+    _shift_prev,
+    _stick_fraction,
+    _type_repr,
     _wandb_run_config,
     train,
 )
@@ -67,6 +85,122 @@ def test_wandb_run_config_handles_missing_model_config():
     assert wcfg["model_config"] == {}
 
 
+# --- AR helper functions (v0.3.0 step 5, docs/v0.3.0-design.md §6) ---------
+
+
+def test_stick_fraction_matches_sigmoid_of_logit():
+    sec_cont = torch.zeros(2, 3, SEC_SLOT_DIM)
+    sec_cont[..., 0] = torch.tensor([[0.0, 2.0, -2.0], [1.0, -1.0, 0.0]])
+    frac = _stick_fraction(sec_cont)
+    assert torch.allclose(frac, torch.sigmoid(sec_cont[..., 0]))
+
+
+def test_remaining_energy_fraction_hand_computed():
+    fraction = torch.tensor([[0.5, 0.5, 1.0]])
+    remaining = _remaining_energy_fraction(fraction)
+    assert torch.allclose(remaining, torch.tensor([[1.0, 0.5, 0.25]]))
+
+
+def test_shift_prev_shifts_and_zero_pads_slot0():
+    x = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float()
+    shifted = _shift_prev(x)
+    assert torch.all(shifted[:, 0] == 0)
+    assert torch.equal(shifted[:, 1:], x[:, :-1])
+
+
+def test_ar_has_prev_false_only_at_slot_zero():
+    has_prev = _ar_has_prev(5, torch.device("cpu"))
+    assert has_prev.shape == (1, 5)
+    assert has_prev.tolist() == [[False, True, True, True, True]]
+
+
+@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
+def test_type_repr_shapes_and_values(target):
+    B, K, emb_dim = 3, 4, 6
+    sec_cont = torch.randn(B, K, SEC_SLOT_DIM)
+    sec_type_idx = torch.randint(0, emb_dim, (B, K))
+    cond_enc = torch.nn.Module()
+    if target == "embedding":
+        cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
+    repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim)
+    expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+    assert repr_.shape == (B, K, expected_width)
+    if target == "physical":
+        assert torch.equal(
+            repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
+        )
+    if target == "onehot":
+        assert torch.all(repr_.sum(-1) == 1.0)
+
+
+@pytest.mark.parametrize(
+    "target,generator",
+    [
+        ("physical", "flow"),
+        ("physical", "wgan"),
+        ("onehot", "flow"),
+        ("onehot", "wgan"),
+        ("embedding", "flow"),
+        ("embedding", "wgan"),
+    ],
+)
+def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
+    target, generator
+):
+    """Regression test tying the refactor together: _assemble_stage2_real is
+    now defined as _assemble_stage2_ar_target(...).flatten(1)."""
+    B, emb_dim = 4, 6
+    sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
+    sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
+    cond_enc = torch.nn.Module()
+    if target == "embedding":
+        cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
+    particle_type_cfg = {"target": target}
+    flat = _assemble_stage2_real(
+        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
+    )
+    unflat = _assemble_stage2_ar_target(
+        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
+    )
+    assert torch.equal(unflat.flatten(1), flat)
+
+
+def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
+    B, emb_dim = 3, 6
+    sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
+    sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
+    cond_enc = torch.nn.Module()
+    out = _assemble_stage2_ar_inputs(
+        sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim
+    )
+    assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
+    assert out["has_prev"].shape == (B, K_MAX)
+    assert out["remaining_frac"].shape == (B, K_MAX)
+    assert out["slot_idx"].shape == (B, K_MAX)
+    assert torch.all(out["slot_idx"][:, 0] == 0.0)
+    assert torch.all(out["slot_idx"][:, -1] == 1.0)
+
+
+def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
+    B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
+    x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
+    grad_probe: dict[str, float] = {}
+    out = _relax_onehot_type_slice(
+        x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe
+    )
+    out.sum().backward()
+    assert grad_probe["cont"] >= 0.0
+    assert grad_probe["type"] >= 0.0
+
+
+def test_relax_onehot_type_slice_grad_probe_none_is_backward_compatible():
+    B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
+    x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
+    out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5)
+    out.sum().backward()
+    assert x_flat.grad is not None
+
+
 # --- end-to-end train() integration tests -----------------------------------
 
 PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
@@ -273,6 +407,59 @@ def _run_train(cfg, out_dir, resume_path=None):
                 ),
             ),
         ),
+        (
+            "ar_wgan_onehot",
+            lambda cfg: (
+                cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+                cfg["stage2_model"].__setitem__(
+                    "particle_type", {"target": "onehot", "lambda": 1.0}
+                ),
+            ),
+        ),
+        (
+            "ar_wgan_physical",
+            lambda cfg: cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+        ),
+        (
+            "ar_flow_onehot",
+            lambda cfg: (
+                cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+                cfg["stage2_model"].__setitem__("generator", "flow"),
+                cfg["stage2_model"].__setitem__(
+                    "particle_type", {"target": "onehot", "lambda": 1.0}
+                ),
+            ),
+        ),
+        (
+            "ar_flow_embedding",
+            lambda cfg: (
+                cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
+                cfg["conditioning"]["material"].__setitem__("type", "embedding"),
+                cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+                cfg["stage2_model"].__setitem__("generator", "flow"),
+                cfg["stage2_model"].__setitem__(
+                    "particle_type", {"target": "embedding", "lambda": 1.0}
+                ),
+            ),
+        ),
+        (
+            "ar_stage2_only",
+            lambda cfg: (
+                cfg["stage1_model"].__setitem__("active", False),
+                cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+            ),
+        ),
+        (
+            "ar_mixed_stage1_wgan_stage2_flow_onehot",
+            lambda cfg: (
+                cfg["stage1_model"].__setitem__("generator", "wgan"),
+                cfg["stage2_model"].__setitem__("generator", "flow"),
+                cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
+                cfg["stage2_model"].__setitem__(
+                    "particle_type", {"target": "onehot", "lambda": 1.0}
+                ),
+            ),
+        ),
     ],
 )
 def test_train_end_to_end(label, mutate):
@@ -400,3 +587,66 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
             ddpm_n_steps=50,
             device=torch.device("cpu"),
         )
+
+
+# --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
+
+
+def test_build_stage_trainers_rejects_scheduled_teacher_forcing():
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["autoregressive"] = {
+        "history": "markov",
+        "teacher_forcing": "scheduled",
+    }
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    with pytest.raises(NotImplementedError):
+        _build_stage_trainers(
+            cfg, models, critics, torch.device("cpu"), total_train_batches=4
+        )
+
+
+def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():
+    """§11.4 differentiability validation-obligation instrumentation: the
+    trunk-gradient-norm-by-slice columns must appear and actually fire for
+    generator='wgan' + particle_type.target='onehot' under decoder=
+    'autoregressive' (added at v0.3.0 step 5 per the design doc's
+    instruction to accrue evidence during the architecture comparison)."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert "stage2_train_grad_norm_type_slice" in rows[0]
+        assert "stage2_train_grad_norm_cont_slice" in rows[0]
+        assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows)
+        assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows)
+
+
+def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
+    """The instrumentation is decoder-agnostic — one_shot + wgan + onehot
+    must populate the same columns."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows)
+        assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows)
+
+
+def test_wgan_physical_omits_grad_norm_slice_columns():
+    cfg = _base_cfg()  # default stage2_model has no particle_type -> "physical"
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
+        assert "stage2_train_grad_norm_type_slice" not in header
+        assert "stage2_train_grad_norm_cont_slice" not in header