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>

giant/config.py

@@ -733,6 +733,22 @@ def validate_config(cfg: dict) -> None:
             "(standalone stage-2 evaluation only, never for rollout)"
         )
 
+    if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
+        history = _get_path(cfg, "stage2_model.autoregressive.history")
+        if history != "markov":
+            raise ValueError(
+                f"stage2_model.autoregressive.history = {history!r} is "
+                "accepted by the schema but not implemented until v0.3.0 "
+                "step 7 — use 'markov'"
+            )
+        teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
+        if teacher_forcing != "always":
+            raise ValueError(
+                "stage2_model.autoregressive.teacher_forcing = "
+                f"{teacher_forcing!r} is accepted by the schema but not "
+                "implemented until v0.3.0 step 7 — use 'always'"
+            )
+
 
 _CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"}
 

giant/model/network.py

@@ -782,6 +782,43 @@ def build_trunk(
     return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
 
 
+# ---------------------------------------------------------------------------
+# History encoders — stage-2 autoregressive only (docs/v0.3.0-design.md §6.2)
+# ---------------------------------------------------------------------------
+
+
+class HistoryEncoder(nn.Module):
+    """Interface for stage-2 autoregressive per-token history summaries:
+    `forward(feat, has_prev) -> (B, K, out_dim)`. `MarkovHistory` is the only
+    implementation until v0.3.0 step 7 (`AttentionHistory`,
+    `history = "attention"`)."""
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class MarkovHistory(HistoryEncoder):
+    """Summarizes the previous secondary's own `(energy_fraction, direction,
+    type_representation)` through one small MLP — the "markov" history
+    (docs/v0.3.0-design.md §6.2): token i+1 only ever sees token i plus the
+    running scalars (`remaining_frac`/`slot_idx`, fused in separately by
+    `Stage2Autoregressive._token_cond`), not the full prefix.
+
+    At slot 0 (`has_prev` False) substitutes a learned start vector rather
+    than zeros — a reasonable default, not specified by the design doc.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU())
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.mlp(x)
+
+
 # ---------------------------------------------------------------------------
 # Stage models (docs/v0.3.0-design.md §5.3)
 # ---------------------------------------------------------------------------
@@ -1041,16 +1078,214 @@ class Stage2OneShot(nn.Module):
 
 
 class Stage2Autoregressive(nn.Module):
-    """Not implemented until v0.3.0 steps 4-7 (docs/v0.3.0-design.md §6, §12)
-    — stub so `stage2_model.decoder = "autoregressive"` (the v0.3.0 default)
-    fails loudly instead of silently no-op-ing."""
+    """Emits secondaries one at a time in descending-energy order
+    (docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous
+    k_max-slot prediction. Only `history = "markov"` is implemented (v0.3.0
+    step 5) — `history = "attention"` raises immediately at construction.
+    `teacher_forcing` handling lives entirely in the trainer
+    (`giant/train.py`), since it only affects how training inputs are
+    assembled, not this module's architecture.
 
-    def __init__(self, *args, **kwargs) -> None:
+    Under teacher forcing every token's conditioning is built from ground
+    truth, so a whole K-token sequence trains in one parallel batched pass:
+    `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
+    to `k_max`) — this also means a future one-token-at-a-time inference loop
+    (`K=1` per call, step 6) needs no interface change here.
+
+    Two independent conditioning paths, mirroring `Stage2OneShot`'s
+    `_cond_embed` but split in two: `_base_cond` (`cond_enc` +
+    `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
+    on token position; `_token_cond` additionally fuses in the history
+    encoding and two running scalars (remaining energy-budget fraction,
+    normalized slot index), and feeds `forward`/`predict_type`/the trunk.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        context_dim: int = 64,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "wgan",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        k_max: int = K_MAX,
+        router: Router | None = None,
+        build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
+        history: str = "markov",
+    ) -> None:
         super().__init__()
+        if history != "markov":
             raise NotImplementedError(
-            "stage2_model.decoder = 'autoregressive' is not implemented yet "
-            "(design doc v0.3.0 steps 4-7) — use decoder = 'one_shot' for now"
+                f"stage2_model.autoregressive.history={history!r} is not "
+                "implemented until v0.3.0 step 7 — use 'markov'"
             )
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        emb_dim = particle_cfg["emb_dim"]
+        self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
+
+        self.cond_enc = ConditionEncoder(
+            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        )
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.base_fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        # Reuses conditioning.out_dim for the history encoder's own output
+        # width — there's no dedicated stage2_model.autoregressive key for
+        # this, a reasonable default rather than a design-doc-specified value.
+        history_dim = cond_out_dim
+        self.history_encoder = MarkovHistory(CONT_SLOT_DIM + self.type_dim, history_dim)
+        token_fuse_in = (
+            cond_out_dim + context_dim + history_dim + 2
+        )  # +2: remaining_frac, slot_idx
+        self.token_fuse = nn.Sequential(
+            nn.Linear(token_fuse_in, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        has_time = generator in ("flow", "ddpm")
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
+        in_dim = noise_dim if generator == "wgan" else token_dim
+        self.trunk = build_trunk(
+            router,
+            in_dim,
+            token_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+        )
+
+        self.n_sec_head = None
+        if build_n_sec_head:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max + 1),
+            )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, self.type_dim),
+            )
+
+    def _base_cond(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
+    ) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.base_fuse(torch.cat([base, ctx], dim=-1))
+
+    def _token_cond(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+    ) -> torch.Tensor:
+        K = history_feat.size(1)
+        base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
+        ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
+        hist = self.history_encoder(history_feat, has_prev)
+        scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
+        return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, K = x_t.shape[0], x_t.shape[1]
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+        )
+        if self.time_emb is not None:
+            assert t is not None
+            t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
+            cond = torch.cat([t_emb, c_emb], dim=-1)
+        else:
+            cond = c_emb
+        x_flat = x_t.reshape(B * K, -1)
+        cond_flat = cond.reshape(B * K, -1)
+        cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
+        return out.view(B, K, -1)
+
+    def predict_n_sec(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
+    ) -> torch.Tensor:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive has no n_sec_head — it belongs to "
+                "a migrated v0.2 checkpoint (legacy_owner='stage1'); call "
+                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
+            )
+        return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
+
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+        )
+        B, K, _ = c_emb.shape
+        return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
 
 
 class CriticModel(nn.Module):
@@ -1299,9 +1534,6 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
 
     if s2cfg.get("active", True):
         decoder = s2cfg.get("decoder", "one_shot")
-        if decoder == "autoregressive":
-            result["stage2"] = Stage2Autoregressive()
-            return result
         router_cfg = s2cfg.get("router") or {}
         stage2_router: Router | None = None
         if router_cfg.get("enabled"):
@@ -1316,6 +1548,29 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
         legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
         k_max = s2cfg.get("k_max", K_MAX)
         particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
+
+        if decoder == "autoregressive":
+            ar_cfg = s2cfg.get("autoregressive") or {}
+            result["stage2"] = Stage2Autoregressive(
+                pdg_vocab=pdg_vocab,
+                mat_vocab=mat_vocab,
+                particle_cfg=particle_cfg,
+                material_cfg=material_cfg,
+                hidden_dim=s2cfg.get("hidden_dim", 256),
+                n_res_blocks=s2cfg.get("n_res_blocks", 6),
+                cond_out_dim=cond_out_dim,
+                context_dim=s2cfg.get("context_dim", 64),
+                dropout=s2cfg.get("dropout", 0.0),
+                generator=generator,
+                time_dim=gen_sub.get("time_dim", 64),
+                noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64),
+                k_max=k_max,
+                router=stage2_router,
+                build_n_sec_head=legacy_owner != "stage1",
+                particle_type_cfg=particle_type_cfg,
+                history=ar_cfg.get("history", "markov"),
+            )
+        else:
             sec_dim = stage2_trunk_sec_dim(
                 particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]
             )
@@ -1377,11 +1632,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             stage="stage1",
         )
 
-    if (
-        s2cfg.get("active", True)
-        and s2cfg.get("generator") == "wgan"
-        and s2cfg.get("decoder", "one_shot") != "autoregressive"
-    ):
+    if s2cfg.get("active", True) and s2cfg.get("generator") == "wgan":
         k_max = s2cfg.get("k_max", K_MAX)
         particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
         in_dim = stage2_trunk_sec_dim(

giant/model/schedule.py

@@ -132,3 +132,70 @@ def flow_matching_loss_secondary(
     phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
     phys_loss = (phys_err * mask).sum() / denom
     return cont_loss + phys_loss
+
+
+def flow_matching_loss_secondary_ar(
+    model: torch.nn.Module,
+    x1: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    history_feat: torch.Tensor,
+    has_prev: torch.Tensor,
+    remaining_frac: torch.Tensor,
+    slot_idx: torch.Tensor,
+    sec_mask: torch.Tensor,
+    type_dim: int | None = None,
+) -> torch.Tensor:
+    """`Stage2Autoregressive` analogue of `flow_matching_loss_secondary`, same
+    masked, per-block (continuous vs. type) loss recipe — but native to
+    `Stage2Autoregressive`'s `(B, K_MAX, token_dim)` I/O and its extra
+    per-token conditioning args, rather than a flattened `(B, K_MAX*token_dim)`
+    vector. Kept as a sibling rather than unified with the flat version: the
+    model call signature differs enough (four extra per-token conditioning
+    tensors) that merging would need an awkward shape-flag + closure.
+
+    Under teacher forcing (docs/v0.3.0-design.md §6.2 point 3) this is still a
+    single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are
+    already built from ground truth for every slot by the caller
+    (`giant.train._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).
+
+    x1: (B, K_MAX, CONT_SLOT_DIM + type_dim) — per-token flattened target
+        (stick_logit, dir, then a `type_dim`-wide type slice)
+    sec_mask: (B, K_MAX) bool — True for valid secondary slots
+    type_dim: as `flow_matching_loss_secondary` — defaults to
+        `PARTICLE_PHYS_DIM`, `0` means no type slice is in `x1` at all.
+    """
+    from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
+
+    if type_dim is None:
+        type_dim = PARTICLE_PHYS_DIM
+
+    B, K, _ = x1.shape
+    t = torch.rand(B, K, device=x1.device)
+    x0 = torch.randn_like(x1)
+    x_t = (1.0 - t.unsqueeze(-1)) * x0 + t.unsqueeze(-1) * x1
+    u_t = x1 - x0
+    v_t = model(
+        x_t,
+        cond_cont,
+        cond_cat,
+        stage1_out,
+        history_feat,
+        has_prev,
+        remaining_frac,
+        slot_idx,
+        t=t,
+    )
+
+    err = (v_t - u_t) ** 2
+    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)
+
+    mask = sec_mask.float()
+    denom = mask.sum().clamp(min=1)
+    cont_loss = (cont_err * mask).sum() / denom
+    if type_dim == 0:
+        return cont_loss
+    phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
+    phys_loss = (phys_err * mask).sum() / denom
+    return cont_loss + phys_loss

giant/train.py

@@ -16,7 +16,7 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from giant.constants import CONT_SLOT_DIM, K_MAX
+from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_to_json
 from giant.model.network import Router, stage2_type_dim
@@ -24,6 +24,7 @@ from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
     flow_matching_loss_secondary,
+    flow_matching_loss_secondary_ar,
 )
 from giant.model.wgan import gradient_penalty, generator_loss
 from giant.validate import validate_marginals
@@ -102,6 +103,67 @@ def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
     return tuple(t.to(device) for t in batch)
 
 
+def _type_repr(
+    sec_type_idx: torch.Tensor,
+    sec_cont: torch.Tensor,
+    particle_type_cfg: dict,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """(B, K_MAX, type_dim) ground-truth type representation, generator-
+    independent (unlike `_assemble_stage2_ar_target`'s training *target*,
+    which varies by generator/objective — see its docstring): `"physical"` ->
+    `(log_mass, charge)`; `"onehot"` -> one-hot of the true class;
+    `"embedding"` -> the conditioning's own detached embedding-table row.
+
+    Used both to build `_assemble_stage2_ar_target`'s wgan+onehot/embedding
+    branch and as the AR history features' previous-secondary identity — the
+    latter must always reflect the true physical secondary that came before,
+    regardless of what the *current* token's own training objective is.
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
+    if target == "onehot":
+        return F.one_hot(sec_type_idx, num_classes=emb_dim).float()
+    return cond_enc.pdg_emb(sec_type_idx).detach()
+
+
+def _assemble_stage2_ar_target(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: dict,
+    generator: str,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """(B, K_MAX, token_dim) ground-truth per-token target — the unflattened
+    analogue of `_assemble_stage2_real` (defined below in terms of this),
+    matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s)
+    own trunk produces for this (target, generator) combination
+    (`giant.model.network.stage2_trunk_sec_dim`; docs/v0.3.0-design.md
+    decision 2/3):
+
+    - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
+      dir, log_mass, charge) as-is.
+    - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
+      just the continuous stick/dir slots — the type slice isn't part of
+      this tensor at all (`type_head` handles it separately).
+    - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
+      slots concatenated with the per-slot type representation (a one-hot of
+      the true class, relaxed on the *generated* side only, by the caller;
+      or the conditioning's own detached embedding-table row).
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return sec_cont
+    cont = sec_cont[..., :CONT_SLOT_DIM]
+    if generator != "wgan":
+        return cont
+    type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
+    return torch.cat([cont, type_repr], dim=-1)
+
+
 def _assemble_stage2_real(
     sec_cont: torch.Tensor,
     sec_type_idx: torch.Tensor,
@@ -110,38 +172,87 @@ def _assemble_stage2_real(
     cond_enc: torch.nn.Module,
     emb_dim: int,
 ) -> torch.Tensor:
-    """Ground-truth flattened stage-2 vector, matching whatever width
-    `Stage2OneShot`'s own trunk produces for this (target, generator)
-    combination (`giant.model.network.stage2_trunk_sec_dim`;
-    docs/v0.3.0-design.md decision 2/3):
+    """Ground-truth flattened stage-2 vector for `Stage2OneShot` — the
+    flattened form of `_assemble_stage2_ar_target`, which
+    `Stage2Autoregressive`'s per-token target also uses; the two must stay in
+    lockstep. See `_assemble_stage2_ar_target`'s docstring for the
+    (target, generator) width rules."""
+    return _assemble_stage2_ar_target(
+        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
+    ).flatten(1)
 
-    - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
-      dir, log_mass, charge) flattened as-is.
-    - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
-      just the continuous stick/dir slots — the type slice isn't part of
-      this vector at all (`Stage2OneShot.type_head` handles it separately).
-    - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
-      slots concatenated with the per-slot type vector — a one-hot of the
-      true class (`"onehot"`, relaxed on the *generated* side only, by the
-      caller) or the conditioning's own detached embedding-table row
-      (`"embedding"`, already continuous — no relaxation needed either side,
-      §2.1).
-    """
-    target = particle_type_cfg.get("target", "physical")
-    if target == "physical":
-        return sec_cont.flatten(1)
-    cont = sec_cont[..., :CONT_SLOT_DIM]
-    if generator != "wgan":
-        return cont.flatten(1)
-    if target == "onehot":
-        type_vec = F.one_hot(sec_type_idx, num_classes=emb_dim).float()
-    else:
-        type_vec = cond_enc.pdg_emb(sec_type_idx).detach()
-    return torch.cat([cont, type_vec], dim=-1).flatten(1)
+
+def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
+    """(B, K_MAX) — sigmoid of each slot's own stick-breaking logit
+    (`sec_cont[...,0]`); scale-free (see `giant.data.transforms.
+    encode_secondaries`), so this needs no absolute `e_sec`."""
+    return torch.sigmoid(sec_cont[..., 0])
+
+
+def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
+    """(B, K_MAX) — fraction of the original e_sec budget unclaimed entering
+    slot i: `1.0` at `i=0`, `prod_{j<i}(1-fraction_j)` for `i>=1`
+    (docs/v0.3.0-design.md §6.3 — "no re-derivation needed": the existing
+    stick-breaking encoding is already scale-free, so this is derivable from
+    the batch's ground-truth stick logits alone, no `e_sec` required)."""
+    cumprod = torch.cumprod(1.0 - fraction, dim=1)
+    return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
+
+
+def _shift_prev(x: torch.Tensor) -> torch.Tensor:
+    """`(B, K, ...)` -> same shape, slot i holds slot i-1's value; slot 0 gets
+    an arbitrary zero placeholder (never read as-is — see `_ar_has_prev`;
+    `MarkovHistory` substitutes its own learned start vector there instead)."""
+    return torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
+
+
+def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
+    """`(1, K_MAX)` bool: True for slot index `>= 1`. Correct without
+    `n_sec`: `sec_mask` is a prefix mask, so any *valid* token at `k>=1`
+    always has a valid predecessor at `k-1`; the only wrong cases are tokens
+    that are themselves padding, already masked out of every loss."""
+    return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
+
+
+def _assemble_stage2_ar_inputs(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: dict,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> dict[str, torch.Tensor]:
+    """Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)`
+    or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means
+    every token's input is ground truth, docs/v0.3.0-design.md §6.2 point 3).
+    Keys match `Stage2Autoregressive.forward`'s trailing kwargs."""
+    device = sec_cont.device
+    B, K = sec_cont.shape[0], sec_cont.shape[1]
+    fraction = _stick_fraction(sec_cont)
+    type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
+    history_feat = torch.cat(
+        [
+            _shift_prev(fraction).unsqueeze(-1),
+            _shift_prev(sec_cont[..., 1:CONT_SLOT_DIM]),
+            _shift_prev(type_repr),
+        ],
+        dim=-1,
+    )
+    slot_idx = (torch.arange(K, device=device).float() / max(K - 1, 1)).unsqueeze(0)
+    return {
+        "history_feat": history_feat,
+        "has_prev": _ar_has_prev(K, device).expand(B, -1),
+        "remaining_frac": _remaining_energy_fraction(fraction),
+        "slot_idx": slot_idx.expand(B, -1),
+    }
 
 
 def _relax_onehot_type_slice(
-    x_flat: torch.Tensor, k_max: int, cont_dim: int, type_dim: int, tau: float
+    x_flat: torch.Tensor,
+    k_max: int,
+    cont_dim: int,
+    type_dim: int,
+    tau: float,
+    grad_probe: dict[str, float] | None = None,
 ) -> torch.Tensor:
     """Straight-through Gumbel-softmax relaxation of the per-slot type slice
     inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator
@@ -149,10 +260,26 @@ def _relax_onehot_type_slice(
     hard one-hot (matching what the critic sees from real data), the
     backward pass flows smooth gradient. Continuous slots (stick/dir, and
     the type slice itself under `target = "embedding"`, which never calls
-    this) pass through unchanged."""
+    this) pass through unchanged.
+
+    `grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2
+    norm of the gradient reaching this split point during the next
+    `.backward()` call that touches it — a backward hook, not a second
+    backward pass. This is the §11.4 differentiability validation-obligation
+    instrumentation (docs/v0.3.0-design.md): the trunk-gradient contribution
+    from the type slice vs. the continuous slices, for
+    `particle_type.target="onehot"` + `generator="wgan"`. Only ever populated
+    on a `did_g_step` batch — the critic step backprops through
+    `fake.detach()`, which never reaches these hooks — so it stays empty
+    (callers default to `0.0`) otherwise."""
     B = x_flat.size(0)
     x = x_flat.view(B, k_max, cont_dim + type_dim)
     cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
+    if grad_probe is not None:
+        cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
+        type_logits.register_hook(
+            lambda g: grad_probe.__setitem__("type", g.norm().item())
+        )
     type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
     return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
 
@@ -234,6 +361,7 @@ class FlowDDPMStageTrainer(StageTrainer):
         device: torch.device,
         particle_type_cfg: dict | None = None,
         particle_type_emb_dim: int = 16,
+        decoder: str = "one_shot",
     ) -> None:
         if is_stage2 and generator not in ("flow",):
             raise NotImplementedError(
@@ -244,6 +372,7 @@ class FlowDDPMStageTrainer(StageTrainer):
         self.name = name
         self.is_stage2 = is_stage2
         self.generator = generator
+        self.decoder = decoder
         self.device = device
         self.model = model.to(device)
         self.lambda_weight = lambda_weight
@@ -302,12 +431,29 @@ class FlowDDPMStageTrainer(StageTrainer):
             return self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx)
         return self.model.predict_n_sec(cond_cont, cond_cat)
 
-    def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx):
+    def _generator_loss(
+        self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None
+    ):
         if not self.is_stage2:
             if self.generator == "flow":
                 return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
             assert self.ddpm_schedule is not None
             return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
+        if self.decoder == "autoregressive":
+            assert ar_inputs is not None
+            return flow_matching_loss_secondary_ar(
+                self.model,
+                x1_s2,
+                cond_cont,
+                cond_cat,
+                stage1_ctx,
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
+                sec_mask,
+                type_dim=self._flow_type_dim,
+            )
         return flow_matching_loss_secondary(
             self.model,
             x1_s2,
@@ -319,10 +465,17 @@ class FlowDDPMStageTrainer(StageTrainer):
         )
 
     def _type_loss(
-        self, cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device
+        self,
+        cond_cont,
+        cond_cat,
+        stage1_ctx,
+        sec_type_idx,
+        sec_mask,
+        device,
+        ar_inputs=None,
     ):
-        """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for
-        `Stage2OneShot.type_head` — the non-adversarial counterpart to
+        """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the
+        stage-2 model's `type_head` — the non-adversarial counterpart to
         WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5).
         Zero when this stage has no `type_head` (stage 1, or
         `particle_type.target = "physical"`)."""
@@ -331,6 +484,18 @@ class FlowDDPMStageTrainer(StageTrainer):
         type_head = getattr(self.model, "type_head", None)
         if not self.is_stage2 or type_head is None:
             return l_type, type_acc
+        if self.decoder == "autoregressive":
+            assert ar_inputs is not None
+            type_out = self.model.predict_type(
+                cond_cont,
+                cond_cat,
+                stage1_ctx,
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
+            )
+        else:
             type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
@@ -360,8 +525,18 @@ class FlowDDPMStageTrainer(StageTrainer):
         ) = _batch_to_device(batch, device)
         sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
         stage1_ctx = x1_s1.detach()
-        x1_s2 = (
-            _assemble_stage2_real(
+
+        x1_s2 = None
+        ar_inputs = None
+        if self.is_stage2 and self.decoder == "autoregressive":
+            ar_inputs = _assemble_stage2_ar_inputs(
+                sec_cont,
+                sec_type_idx,
+                self.particle_type_cfg,
+                self.model.cond_enc,
+                self.particle_type_emb_dim,
+            )
+            x1_s2 = _assemble_stage2_ar_target(
                 sec_cont,
                 sec_type_idx,
                 self.particle_type_cfg,
@@ -369,12 +544,18 @@ class FlowDDPMStageTrainer(StageTrainer):
                 self.model.cond_enc,
                 self.particle_type_emb_dim,
             )
-            if self.is_stage2
-            else None
+        elif self.is_stage2:
+            x1_s2 = _assemble_stage2_real(
+                sec_cont,
+                sec_type_idx,
+                self.particle_type_cfg,
+                self.generator,
+                self.model.cond_enc,
+                self.particle_type_emb_dim,
             )
 
         l_gen = self._generator_loss(
-            cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx
+            cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs
         )
         n_sec_logits = self._predict_n_sec(cond_cont, cond_cat, stage1_ctx)
         l_nsec = torch.zeros((), device=device)
@@ -384,7 +565,13 @@ class FlowDDPMStageTrainer(StageTrainer):
             nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()
 
         l_type, type_acc = self._type_loss(
-            cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device
+            cond_cont,
+            cond_cat,
+            stage1_ctx,
+            sec_type_idx,
+            sec_mask,
+            device,
+            ar_inputs=ar_inputs,
         )
 
         l_balance = l_proc = l_entropy = torch.zeros((), device=device)
@@ -531,9 +718,11 @@ class WGANStageTrainer(StageTrainer):
         particle_type_emb_dim: int = 16,
         type_gumbel_tau_start: float = 1.0,
         type_gumbel_tau_end: float = 0.1,
+        decoder: str = "one_shot",
     ) -> None:
         self.name = name
         self.is_stage2 = is_stage2
+        self.decoder = decoder
         self.device = device
         self.model = model.to(device)
         self.critic = critic.to(device)
@@ -592,6 +781,7 @@ class WGANStageTrainer(StageTrainer):
         ) = _batch_to_device(batch, device)
         B = cond_cont.size(0)
         stage1_ctx = x1_s1.detach()
+        grad_probe: dict[str, float] = {}
 
         if not self.is_stage2:
             real = x1_s1
@@ -615,6 +805,41 @@ class WGANStageTrainer(StageTrainer):
             mask = (
                 sec_mask.unsqueeze(-1).expand(-1, -1, slot_width).reshape(B, -1).float()
             )
+
+            def critic_fn(x):
+                return self.critic(x, cond_cont, cond_cat, stage1_ctx)
+
+            if self.decoder == "autoregressive":
+                ar = _assemble_stage2_ar_inputs(
+                    sec_cont,
+                    sec_type_idx,
+                    self.particle_type_cfg,
+                    self.model.cond_enc,
+                    self.particle_type_emb_dim,
+                )
+                real = (
+                    _assemble_stage2_ar_target(
+                        sec_cont,
+                        sec_type_idx,
+                        self.particle_type_cfg,
+                        "wgan",
+                        self.model.cond_enc,
+                        self.particle_type_emb_dim,
+                    ).reshape(B, -1)
+                    * mask
+                )
+                z = torch.randn(B, K_MAX, self.model.noise_dim, device=device)
+                fake_raw = self.model(
+                    z,
+                    cond_cont,
+                    cond_cat,
+                    stage1_ctx,
+                    ar["history_feat"],
+                    ar["has_prev"],
+                    ar["remaining_frac"],
+                    ar["slot_idx"],
+                ).reshape(B, -1)
+            else:
                 real = (
                     _assemble_stage2_real(
                         sec_cont,
@@ -626,17 +851,16 @@ class WGANStageTrainer(StageTrainer):
                     )
                     * mask
                 )
-
-            def critic_fn(x):
-                return self.critic(x, cond_cont, cond_cat, stage1_ctx)
-
                 z = torch.randn(B, self.model.noise_dim, device=device)
                 fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)
+
             if target == "onehot":
                 # Straight-through Gumbel-softmax relaxation of the type
                 # slice only (decision 5) — the critic must see a hard
                 # one-hot forward (matching what "real" data looks like)
                 # while gradient still flows smoothly to the generator.
+                # grad_probe captures the §11.4 gradient-magnitude
+                # instrumentation — see _relax_onehot_type_slice's docstring.
                 tau = _gumbel_tau(
                     global_step,
                     self.total_steps,
@@ -644,7 +868,7 @@ class WGANStageTrainer(StageTrainer):
                     self.type_gumbel_tau_end,
                 )
                 fake_raw = _relax_onehot_type_slice(
-                    fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau
+                    fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe
                 )
             fake = fake_raw * mask
 
@@ -711,6 +935,8 @@ class WGANStageTrainer(StageTrainer):
             "grad_norm": grad_norm_d.item() + grad_norm_g.item(),
             "grad_norm_d": grad_norm_d.item(),
             "grad_norm_g": grad_norm_g.item(),
+            "grad_norm_type_slice": grad_probe.get("type", 0.0),
+            "grad_norm_cont_slice": grad_probe.get("cont", 0.0),
             "lr": self.optimizer.param_groups[0]["lr"],
             "critic_lr": self.optimizer_d.param_groups[0]["lr"],
             "batch_size": B,
@@ -820,6 +1046,17 @@ def _build_stage_trainers(
             "target": "physical"
         }
         particle_type_emb_dim = cfg["conditioning"]["particle"]["emb_dim"]
+        decoder = stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot"
+        if is_stage2 and decoder == "autoregressive":
+            teacher_forcing = (stage_cfg.get("autoregressive") or {}).get(
+                "teacher_forcing", "always"
+            )
+            if teacher_forcing != "always":
+                raise NotImplementedError(
+                    "stage2_model.autoregressive.teacher_forcing="
+                    f"{teacher_forcing!r} is not implemented until v0.3.0 "
+                    "step 7 — use 'always'"
+                )
 
         if generator == "wgan":
             critic = critics.get(name)
@@ -848,6 +1085,7 @@ def _build_stage_trainers(
                 particle_type_emb_dim=particle_type_emb_dim,
                 type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
                 type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
+                decoder=decoder,
             )
         else:
             ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000)
@@ -873,6 +1111,7 @@ def _build_stage_trainers(
                 device=device,
                 particle_type_cfg=particle_type_cfg,
                 particle_type_emb_dim=particle_type_emb_dim,
+                decoder=decoder,
             )
     return trainers
 
@@ -892,6 +1131,14 @@ def _metrics_fields(trainers: dict[str, StageTrainer]) -> list[str]:
                 f"{name}_train_grad_norm_d",
                 f"{name}_train_grad_norm_g",
             ]
+            if (
+                trainer.is_stage2
+                and trainer.particle_type_cfg.get("target") == "onehot"
+            ):
+                fields += [
+                    f"{name}_train_grad_norm_type_slice",
+                    f"{name}_train_grad_norm_cont_slice",
+                ]
         else:
             fields += [
                 f"{name}_train_loss",
@@ -1358,6 +1605,13 @@ def train(
                     metrics_row[f"{name}_train_grad_norm_g"] = (
                         sums.get("grad_norm_g", 0.0) / n_train
                     )
+                    if tr.is_stage2 and tr.particle_type_cfg.get("target") == "onehot":
+                        metrics_row[f"{name}_train_grad_norm_type_slice"] = (
+                            sums.get("grad_norm_type_slice", 0.0) / n_train
+                        )
+                        metrics_row[f"{name}_train_grad_norm_cont_slice"] = (
+                            sums.get("grad_norm_cont_slice", 0.0) / n_train
+                        )
                     grad_norm_total += sums.get("grad_norm", 0.0) / n_train
                     metrics_row[f"{name}_critic_lr"] = tr.optimizer_d.param_groups[0][
                         "lr"

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