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

@@ -637,6 +637,55 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_ar_default_markov_always_passes():
+    """DEFAULT_CONFIG already has decoder='autoregressive',
+    history='markov', teacher_forcing='always' — must not raise (v0.3.0
+    step 5; see also test_validate_config_default_config_passes)."""
+    cfg = _cfg_with(**{"stage2_model.decoder": "autoregressive"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_ar_history_attention_not_implemented():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.history": "attention",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "history" in str(e)
+
+
+def test_validate_config_ar_teacher_forcing_scheduled_not_implemented():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.teacher_forcing": "scheduled",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "teacher_forcing" in str(e)
+
+
+def test_validate_config_ar_checks_skipped_under_one_shot():
+    """history/teacher_forcing values that would fail under AR are irrelevant
+    (and unchecked) when decoder='one_shot'."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "one_shot",
+            "stage2_model.autoregressive.history": "attention",
+            "stage2_model.autoregressive.teacher_forcing": "scheduled",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 # ---------------------------------------------------------------------------
 # checkpoint config-mismatch warnings (unchanged surface, still exercised)
 # ---------------------------------------------------------------------------

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}"

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 ──────────────────────────────────────────────────────────────────
 
 

tests/test_pipeline.py

@@ -105,10 +105,10 @@ def _tiny_cfg(**train_overrides):
     cfg["train"].update(train_overrides)
     cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0})
     cfg["stage2_model"].update(
-        # decoder="autoregressive" is DEFAULT_CONFIG's default (the finished
-        # v0.3.0 target) but Stage2Autoregressive isn't implemented until
-        # design doc step 4/5 — every run must override to "one_shot" for now.
-        {"decoder": "one_shot", "hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}
+        # decoder="autoregressive" is DEFAULT_CONFIG's default (v0.3.0 step 5)
+        # and left as-is here on purpose, so this pipeline-level fixture
+        # exercises the real default end-to-end against actual data.
+        {"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}
     )
     cfg["conditioning"]["particle"]["emb_dim"] = 4
     cfg["conditioning"]["material"]["emb_dim"] = 4

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