v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies

Works through docs/v0.3.0-followups.md item by item, closing the gap
between the design doc and the shipped v0.3.0-stage2-autoregressive code:

1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2
   dispatch, stage-2 particle-type-class marginal.
2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run.
3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/
   train instead of the hardcoded K_MAX constant.
4. Mixed conditioning.particle.type / conditioning.material.type support
   end-to-end (data pipeline + dwarf warm-cache).
5. conditioning.share_stages = true: one shared ConditionEncoder instance
   across both stages.
6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2
   (was silently unimplemented).
7. giant predict/rollout: implement conditioning.*.type = "onehot" via the
   checkpoint's saved pdg_topn_map/mat_topn_map.
8. network.py's checkpoint-path model_config migration now fails loudly on
   non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's
   TOML-load path (§4.2).
9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a
   rollout-capable checkpoint (§9).

Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics),
mostly a test-helper dict-unpack pattern that made every unrelated
constructor keyword look like a type error.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_cli_predict.py

@@ -1,73 +1,14 @@
 import uuid
 
-import pytest
-import typer
 import yaml
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
-    _check_conditioning_onehot_support,
     _resolve_prediction_output,
     _write_prediction_ref,
 )
 
 
-# ---------------------------------------------------------------------------
-# _check_conditioning_onehot_support
-# ---------------------------------------------------------------------------
-
-
-def _nested_model_cfg(
-    particle_type="physical", material_type="physical", target="physical"
-):
-    return {
-        "conditioning": {
-            "particle": {"type": particle_type, "emb_dim": 8},
-            "material": {"type": material_type, "emb_dim": 8},
-        },
-        "stage2_model": {"particle_type": {"target": target}},
-    }
-
-
-def test_check_conditioning_onehot_support_allows_physical():
-    _check_conditioning_onehot_support(_nested_model_cfg(), "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning():
-    cfg = _nested_model_cfg(particle_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "predict")
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning():
-    cfg = _nested_model_cfg(material_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "rollout")
-
-
-def test_check_conditioning_onehot_support_allows_onehot_particle_type_target():
-    """stage2_model.particle_type.target="onehot" is implemented (v0.3.0
-    step 6, giant.rollout.decode_secondary_identity) — it's a separate axis
-    from conditioning.particle.type, which this guard doesn't gate at all."""
-    cfg = _nested_model_cfg(target="onehot")
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_allows_embedding_particle_type_target():
-    cfg = _nested_model_cfg(
-        particle_type="embedding", material_type="embedding", target="embedding"
-    )
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_is_noop_for_v02_flat_model_config():
-    """A v0.2 checkpoint's flat model_config has conditioning as a plain
-    string, not a dict — never onehot, so this must be a silent no-op rather
-    than crash on `.get("particle")` against a string."""
-    cfg = {"conditioning": "embedding", "mode": "flow"}
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
 # ---------------------------------------------------------------------------

tests/test_cli_train_overrides.py

@@ -0,0 +1,97 @@
+"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
+decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
+independently override each stage's config block, and must take precedence
+over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
+apply the same value to both stages for backward compatibility."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+import giant.cli as cli
+
+runner = CliRunner()
+
+
+def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["cfg"] = cfg
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
+    )
+    assert result.exit_code == 0, result.output
+    return captured["cfg"]
+
+
+def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--mode", "wgan", "--stage1-generator", "flow"],
+    )
+    assert cfg["stage1_model"]["generator"] == "flow"
+    assert cfg["stage2_model"]["generator"] == "wgan"
+
+
+def test_stage2_only_knobs(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--stage2-decoder",
+            "one_shot",
+            "--stage2-k-max",
+            "8",
+            "--stage2-hidden-dim",
+            "32",
+            "--stage2-context-dim",
+            "16",
+            "--stage2-stage1-context",
+            "sampled",
+        ],
+    )
+    assert cfg["stage2_model"]["decoder"] == "one_shot"
+    assert cfg["stage2_model"]["k_max"] == 8
+    assert cfg["stage2_model"]["hidden_dim"] == 32
+    assert cfg["stage2_model"]["context_dim"] == 16
+    assert cfg["stage2_model"]["stage1_context"] == "sampled"
+    # untouched stage1 defaults
+    assert cfg["stage1_model"]["hidden_dim"] == 256
+
+
+def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
+    )
+    assert cfg["stage1_model"]["hidden_dim"] == 128
+
+
+def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--mode",
+            "wgan",
+            "--n-critic",
+            "5",
+            "--stage1-n-critic",
+            "3",
+            "--stage2-gp-weight",
+            "2.5",
+        ],
+    )
+    assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
+    assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
+    assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
+    assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5

tests/test_config.py

@@ -584,21 +584,19 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
     gconfig.validate_config(cfg)  # must not raise
 
 
-def test_validate_config_mixed_particle_material_conditioning_not_supported():
-    """The data pipeline doesn't support mixed conditioning types yet, even
-    though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1
-    vs. giant/data/transforms.py's still-single conditioning param)."""
+def test_validate_config_mixed_particle_material_conditioning_is_valid():
+    """docs/v0.3.0-design.md §3.1: the particle and material conditioning
+    axes are configured independently and may mix freely — e.g. material
+    "physical" with particle "embedding" — and the data pipeline
+    (giant/data/transforms.py) now implements that end-to-end, so
+    validate_config must not reject it."""
     cfg = _cfg_with(
         **{
             "conditioning.particle.type": "physical",
             "conditioning.material.type": "embedding",
         }
     )
-    try:
-        gconfig.validate_config(cfg)
-        assert False, "expected ValueError"
-    except ValueError as e:
-        assert "mixed" in str(e) or "conditioning.material.type" in str(e)
+    gconfig.validate_config(cfg)  # must not raise
 
 
 def test_validate_config_pdg_router_incompatible_with_physical_conditioning():
@@ -639,6 +637,49 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
+    """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a
+    rollout-capable checkpoint' — both stages active means giant rollout
+    could load this checkpoint, but 'truth' has no ground truth to draw
+    n_sec from at rollout time."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": True,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_sec.mode" in str(e) and "truth" in str(e)
+
+
+def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint():
+    """'truth' is exactly the standalone stage-2 evaluation mode the design
+    doc carves out — stage1_model.active = false must still pass."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": False,
+            "stage2_model.active": True,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_n_sec_truth_allowed_when_stage2_inactive():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": False,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 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

tests/test_dataset.py

@@ -126,7 +126,8 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
             target_normalizer=tgt_norm,
             batch_size=4,
             shuffle=False,
-            conditioning="embedding",
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
         )
         return sum(len(batch[0]) for batch in ds)
 

tests/test_dwarf.py

@@ -103,7 +103,7 @@ def test_warm_cache_writes_sidecar(tmp_path):
     assert loaded is not None
     assert loaded.vocab is not None
     assert loaded.event_index is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
 
 
 def test_warm_cache_second_run_hits_cache(tmp_path):
@@ -167,5 +167,5 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
     assert "fitting normalizer (streaming)" in result.output
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
-    assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
+    assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers

tests/test_flow.py

@@ -52,7 +52,7 @@ def test_sample_flow_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 def test_ddpm_loss_nonneg():
@@ -69,4 +69,4 @@ def test_sample_ddim_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)

tests/test_migration_v02_v03.py

@@ -198,6 +198,56 @@ def test_migrate_legacy_model_config_shape():
     assert migrated["stage2_model"]["decoder"] == "one_shot"
 
 
+def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
+    """docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's
+    model_config carrying a non-default expert_hidden_dim/expert_n_blocks
+    must fail loudly through this path too (§4.2) — not just
+    giant.config.migrate_config's parallel TOML-load path. Silently dropping
+    these keys (build_router's kwarg filtering) would resize the experts
+    instead of refusing."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net._migrate_legacy_model_config(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
+def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 0,
+        "expert_n_blocks": 0,
+    }
+    migrated = net._migrate_legacy_model_config(legacy_cfg)
+    assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
+    assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage2_model"]["router"]
+
+
+def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
+    """The same check must also fire through the actual caller,
+    build_models, not just the internal helper directly."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net.build_models(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
 def test_build_models_accepts_new_nested_shape_unchanged():
     """A dict that already has a 'stage1_model' key (the new shape) is
     passed through build_models without going through the legacy migration

tests/test_network.py

@@ -1,5 +1,8 @@
+import copy
+
 import pytest
 import torch
+from giant import config as gconfig
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
     AttentionHistory,
@@ -9,6 +12,7 @@ from giant.model.network import (
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
     stage2_type_dim,
@@ -628,3 +632,52 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
 def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
     model = _build_stage2_ar("physical", "wgan", history="markov")
     assert model.init_history_cache() is None
+
+
+# ── build_models: conditioning.share_stages ─────────────────────────────────
+
+
+def _minimal_model_config(share_stages: bool) -> dict:
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["share_stages"] = share_stages
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    return {
+        "pdg_vocab": 3,
+        "mat_vocab": 2,
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def test_build_models_share_stages_true_shares_condition_encoder_instance():
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+
+def test_build_models_share_stages_false_builds_independent_condition_encoders():
+    built = build_models(_minimal_model_config(share_stages=False))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is not stage2.cond_enc
+
+
+def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
+    """The shared encoder's parameters must actually appear in both stages'
+    own `.parameters()` — that's what makes each stage's independent
+    optimizer include (and update) them, which is the actual mechanism behind
+    "shared weights, forced common representation" (docs/v0.3.0-design.md
+    §3.1), not just object identity on `.cond_enc`."""
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+
+    shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
+    assert shared_ids
+    assert shared_ids <= {id(p) for p in stage1.parameters()}
+    assert shared_ids <= {id(p) for p in stage2.parameters()}

tests/test_pipeline.py

@@ -7,6 +7,7 @@ import torch
 
 from giant import config as gconfig
 from giant.data import setup_cache
+from giant.data.transforms import Normalizer
 from giant.pipeline import run_train_job
 
 
@@ -235,6 +236,74 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
     assert "fitting normalizer (streaming)" in joined
 
 
+def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other
+    than the K_MAX module constant's default (15) must not produce a shape
+    mismatch between the data pipeline (loader.py/transforms.py padding) and
+    the model (network.py's trunks, sized from this same config value)."""
+    cfg = _tiny_cfg()
+    cfg["stage2_model"]["k_max"] = 3
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["stage2_model"]["k_max"] == 3
+
+
+def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type
+    and conditioning.material.type are configured independently and may mix
+    freely (docs/v0.3.0-design.md §3.1) — e.g. particle "embedding" with
+    material "physical" — end-to-end through the real data pipeline, not
+    just accepted by validate_config."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "embedding"
+    cfg["conditioning"]["material"]["type"] = "physical"
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    cond_cfg = ckpt["model_config"]["conditioning"]
+    assert cond_cfg["particle"]["type"] == "embedding"
+    assert cond_cfg["material"]["type"] == "physical"
+
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    # Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays
+    # unfitted (mean=0/std=1) since "embedding" never computes real values
+    # for it; the material block is fit for real under "physical".
+    from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
+
+    assert cond_norm.mean is not None and cond_norm.std is not None
+    np.testing.assert_allclose(
+        cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
+    )
+    np.testing.assert_allclose(
+        cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
+    )
+    material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
+    assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
+
+
+def test_run_train_job_share_stages_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages
+    = true must actually train (not raise NotImplementedError), and the
+    resulting checkpoint's two stages must reload into a single shared
+    ConditionEncoder instance rather than two independent ones."""
+    from giant.model.network import build_models
+
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["share_stages"] = True
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["conditioning"]["share_stages"] is True
+
+    built = build_models(ckpt["model_config"])
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+    stage1.load_state_dict(ckpt["model"])
+    stage2.load_state_dict(ckpt["sec_decoder"])
+    for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()):
+        assert torch.equal(p1, p2)
+
+
 def test_run_train_job_matches_uncached_output(tmp_path, data):
     _run(data, tmp_path / "uncached", cache_setup=False)
     _run(data, tmp_path / "cached1", cache_setup=True)

tests/test_rollout.py

@@ -111,7 +111,8 @@ def _run(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
     )
 
 
@@ -172,7 +173,7 @@ def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup():
     frontier construction — mass/charge are simply zero-filled, unused."""
     seeds = _seeds(3)
     seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
-    fr, _counts = make_seed_frontier(**seeds, conditioning="embedding")
+    fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding")
     np.testing.assert_array_equal(fr["mass"], 0.0)
     np.testing.assert_array_equal(fr["charge"], 0.0)
 
@@ -384,25 +385,43 @@ def _models_v3(
         noise_dim=8,
     )
     particle_type_cfg = {"target": target}
-    common = dict(
-        pdg_vocab=3,
-        mat_vocab=2,
-        particle_cfg=particle_cfg,
-        material_cfg=material_cfg,
-        hidden_dim=32,
-        n_res_blocks=2,
-        generator=generator2,
-        time_dim=16,
-        noise_dim=8,
-        k_max=k_max,
-        particle_type_cfg=particle_type_cfg,
-        build_n_sec_head=stage2_has_n_sec_head,
-    )
+    # Explicit kwargs rather than a shared **common dict: a dict() call whose
+    # values have heterogeneous types (str/int/dict/bool) widens under static
+    # analysis to dict[str, <big union>], which then makes every constructor
+    # keyword not itself part of that union (router, cond_enc, ...) look like
+    # a type mismatch to `ty` even though every actual value passed is fine.
     if decoder == "one_shot":
         sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim)
-        s2 = Stage2OneShot(sec_dim=sec_dim, **common)
+        s2 = Stage2OneShot(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+            sec_dim=sec_dim,
+        )
     else:
-        s2 = Stage2Autoregressive(**common)
+        s2 = Stage2Autoregressive(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+        )
     return s1.eval(), s2.eval()
 
 
@@ -440,7 +459,8 @@ def _run_v3(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
         pdg_topn_map=pdg_topn_map,
         other_policy=other_policy,
         seed=seed,
@@ -535,6 +555,88 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
         _run_v3(s1, s2, pdg_topn_map=None)
 
 
+# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md
+# item 7) — a separate axis from stage2_model.particle_type.target above: this
+# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own
+# "onehot" mode, not the secondary-species decode. ---------------------------
+
+COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
+COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})
+
+
+def _onehot_conditioning_models():
+    particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    )
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3),
+        generator="flow",
+        time_dim=16,
+    )
+    return s1.eval(), s2.eval()
+
+
+def _run_onehot_conditioning(
+    pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
+):
+    s1, s2 = _onehot_conditioning_models()
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=30,
+        steps=4,
+        batch_size=128,
+        max_tracks_per_event=300,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+    )
+
+
+def test_rollout_conditioning_onehot_end_to_end(fake_material_props):
+    rec = _run_onehot_conditioning()
+    assert len(rec["event_id"]) > 0
+    assert set(rec["event_id"].tolist()) == set(range(6))
+
+
+def test_rollout_conditioning_onehot_particle_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="pdg_topn_map"):
+        _run_onehot_conditioning(pdg_topn_map=None)
+
+
+def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="mat_topn_map"):
+        _run_onehot_conditioning(mat_topn_map=None)
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the

tests/test_sample.py

@@ -150,7 +150,7 @@ def test_sample_flow_returns_n_sec_for_legacy_stage1():
     )
     cond_cont, cond_cat = _cond(5)
     _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
-    assert n_sec.shape == (5,)
+    assert n_sec is not None and n_sec.shape == (5,)
 
 
 # ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────

tests/test_transforms.py

@@ -426,7 +426,13 @@ def test_build_features_embedding_mode_zero_fills_physical_columns():
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
@@ -438,7 +444,13 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     mass, charge = particle_mass_charge(11)
@@ -461,7 +473,13 @@ def test_build_features_physical_mode_unfilled_material_raises():
     pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}
 
     with pytest.raises(MaterialPropertiesNotFilledError):
-        build_features(data, pdg_map, mat_map, conditioning="physical")
+        build_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="physical",
+            material_conditioning="physical",
+        )
 
 
 def test_build_cond_features_mass_charge_override(fake_material_props):
@@ -474,7 +492,13 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
     data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, _ = build_cond_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     np.testing.assert_allclose(
         cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
@@ -494,7 +518,12 @@ def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
     legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
 
     cond_cont, _ = build_cond_features(
-        data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
+        data,
+        pdg_map,
+        mat_map,
+        cond_normalizer=legacy_norm,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
     )
 
     assert cond_cont.shape[-1] == COND_DIM
@@ -521,7 +550,8 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
             pdg_map,
             mat_map,
             cond_normalizer=legacy_norm,
-            conditioning="physical",
+            particle_conditioning="physical",
+            material_conditioning="physical",
         )
 
 
@@ -610,13 +640,23 @@ def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_materi
     }
 
     cond_cont, cond_cat = build_cond_features(
-        data, pdg_map, mat_map, conditioning="physical"
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
     )
     assert cond_cont.shape[-1] == COND_DIM
     np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise
 
     with pytest.raises(KeyError):
-        build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
+        build_cond_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
+        )
 
 
 # ── _WelfordAccumulator ──────────────────────────────────────────────────────

tests/test_validate.py

@@ -1,15 +1,18 @@
 import numpy as np
 import torch
 
-from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
-from giant.model.network import Stage1Model, Stage2OneShot
+from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
 from giant.validate import validate_marginals
 
 _PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 _MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+_K_MAX = 5
 
 
-def _tiny_models():
+def _tiny_models(particle_type_cfg: dict | None = None):
+    """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so
+    n_sec always comes from Stage2OneShot."""
     s1 = Stage1Model(
         pdg_vocab=3,
         mat_vocab=2,
@@ -17,7 +20,13 @@ def _tiny_models():
         material_cfg=_MATERIAL_CFG,
         hidden_dim=16,
         n_res_blocks=1,
-        n_sec_head_k_max=K_MAX,
+    )
+    target = (particle_type_cfg or {}).get("target", "physical")
+    sec_dim = stage2_trunk_sec_dim(
+        particle_type_cfg or {"target": "physical"},
+        "flow",
+        _K_MAX,
+        int(_PARTICLE_CFG["emb_dim"]),
     )
     s2 = Stage2OneShot(
         pdg_vocab=3,
@@ -28,23 +37,29 @@ def _tiny_models():
         n_res_blocks=1,
         generator="flow",
         time_dim=16,
+        k_max=_K_MAX,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
     )
+    assert s2.particle_type_cfg.get("target", "physical") == target
     return s1.eval(), s2.eval()
 
 
-def _zero_secondaries_loader(B=4, n_batches=2):
-    """A val_loader whose every batch has n_sec=0 (real side) — matches the
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape
-    StreamingStepsDataset yields."""
+def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
+    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
+    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
         cond_cat = torch.zeros(B, 2, dtype=torch.long)
         x1 = torch.randn(B, X_DIM)
-        n_sec = torch.zeros(B, dtype=torch.long)
-        sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM)
+        n_sec = torch.full((B,), n_sec_value, dtype=torch.long)
+        sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
+        sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
+        batches.append(
+            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
+        )
     return batches
 
 
@@ -52,22 +67,56 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
     """If n_sec_pred collapses to 0 across the whole validated set (realistic
     during early/unstable training), phys_kl must degrade to NaN instead of
     crashing on the empty-array .min()/.max() reduction inside
-    _histogram_kl -- a regression the old species/bincount code this
-    replaced explicitly guarded against."""
+    _histogram_kl."""
     s1, s2 = _tiny_models()
-    loader = _zero_secondaries_loader()
+    loader = _loader(n_sec_value=0)
 
-    # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so
-    # the generated side's valid-slot mask is also empty (real side is
-    # already all n_sec=0 by construction of the fake loader above).
-    def _fake_sample_flow(model, cond_cont, cond_cat, **kw):
-        B = cond_cont.size(0)
-        return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long)
+    def _fake_resolve_n_sec(
+        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
+    ):
+        return torch.zeros(cond_cont.size(0), dtype=torch.long)
 
-    monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow)
+    monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
 
-    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2)
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2)
 
     assert np.asarray(result["phys_real"]).shape == (0, 2)
     assert np.asarray(result["phys_generated"]).shape == (0, 2)
     assert np.isnan(np.asarray(result["phys_kl"])).all()
+
+
+def test_validate_marginals_physical_target_shapes():
+    s1, s2 = _tiny_models()
+    loader = _loader(n_sec_value=2)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert np.asarray(result["real"]).shape == (4, X_DIM)
+    assert np.asarray(result["generated"]).shape == (4, X_DIM)
+    assert np.asarray(result["kl_divergence"]).shape == (X_DIM,)
+    assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result
+    assert "type_class_real" not in result
+
+
+def test_validate_marginals_onehot_type_class_marginal():
+    particle_type_cfg = {"target": "onehot"}
+    s1, s2 = _tiny_models(particle_type_cfg)
+    loader = _loader(n_sec_value=2, n_classes=s2.type_dim)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert "phys_real" not in result
+    # Real/generated valid-slot counts need not agree (real: ground-truth
+    # n_sec=2 always; generated: the untrained n_sec_head's own prediction).
+    assert np.asarray(result["type_class_real"]).ndim == 1
+    assert np.asarray(result["type_class_gen"]).ndim == 1
+    assert np.asarray(result["type_class_real"]).shape[0] > 0
+
+
+def test_validate_marginals_without_sec_decoder_returns_stage1_only():
+    s1, _ = _tiny_models()
+    loader = _loader(n_sec_value=0)
+
+    result = validate_marginals(s1, loader, n_batches=1, steps=2)
+
+    assert set(result) == {"real", "generated", "kl_divergence"}

tests/test_wgan.py

@@ -120,7 +120,7 @@ def test_sample_wgan_shape():
     cond_cont, cond_cat = _cond(B)
     sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
     assert sample.shape == (B, X_DIM)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 # --- Stage-2 generator/critic ---