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