Merge remote-tracking branch 'origin/master' into fix/issue-83

# Conflicts:
#	.gitea/workflows/ci.yml
#	giant/cli.py

tests/test_analysis_reduce.py

@@ -13,6 +13,7 @@ from giant.analysis.sources import (
     open_side,
     physical_steps,
     secondaries,
+    secondaries_by_step,
 )
 from giant.data.loader import EVENT_ID_FILE_STRIDE
 
@@ -159,18 +160,17 @@ def test_secondaries_rollout_vs_reference_align():
     assert t["pdg"].to_list() == [22, 22]
 
 
-def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
-    r_phys = physical_steps(_rollout_frame(), Side.rollout)
-    r_sec = secondaries(_rollout_frame(), Side.rollout)
-    ev, n = R.sec_count_by_event(r_phys, r_sec)
-    # event 1 has one secondary track; event 2 has none and must still appear (as 0),
-    # not silently drop out of a plain group_by on the secondaries frame alone.
-    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
+def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
+    r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
+    assert r["pdg"].to_list() == [22]
+    # the rollout key is (event_id, parent_id, birth position) — the parent
+    # step's post_pos, copied verbatim onto the child's birth row.
+    assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.0}
 
-    t_all = _reference_frame()
-    t_sec = secondaries(t_all, Side.reference)
-    ev, n = R.sec_count_by_event(t_all, t_sec)
-    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
+    t = secondaries_by_step(_reference_frame(), Side.reference).collect()
+    assert t["pdg"].to_list() == [22, 22]
+    # one row per emitting step; the empty-list step drops out entirely
+    assert [k["_row"] for k in t["step_key"]] == [0, 2]
 
 
 def test_leakage_fraction():

tests/test_catalog.py

@@ -10,10 +10,10 @@ from giant.analysis.catalog import (
     Bundle,
     PlotSpec,
     _containment_depths,
-    _integer_confusion,
     _ks_statistic,
 )
 from giant.analysis.context import Context, build_context
+from giant.analysis.grouping import pdg_label
 from giant.analysis.sources import RolloutSpec
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
@@ -160,8 +160,9 @@ def _validate_payload(r, names: list[str]) -> None:
 # data-dependent edges (event_total_edep), concat-then-mean/std (shower_
 # longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
 # ratio (species_edep_share), a chunkable=False passthrough (router_gating),
-# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
-# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
+# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
+# nested sum-merge into a scorecard (marginal_distance_summary), and
+# concat-then-per-event-derived-quantity
 # (shower_containment_depth_90, reusing the profile matrix's own merge shape).
 _CHUNK_EQUIVALENCE_IDS = [
     "marginal_edep",
@@ -170,9 +171,10 @@ _CHUNK_EQUIVALENCE_IDS = [
     "shower_longitudinal",
     "leakage_fraction",
     "sec_count_per_species",
+    "sec_count_per_step",
+    "sec_count_per_step_by_species",
     "router_gating",
     "marginal_distance_summary",
-    "n_sec_confusion",
     "shower_containment_depth_90",
 ]
 
@@ -217,7 +219,7 @@ def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
 
 
 # ---------------------------------------------------------------------------
-# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
+# new (gitea #76) reductions: KS distance and containment depth
 # ---------------------------------------------------------------------------
 
 
@@ -228,28 +230,6 @@ def test_ks_statistic():
     assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch
 
 
-def test_integer_confusion_matches_event_pairing():
-    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
-    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
-    assert labels == ["0", "1+"]
-    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred
-
-
-def test_integer_confusion_caps_pathological_outliers():
-    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
-    assert labels[-1] == "4+"
-    assert mat.shape == (5, 5)
-    assert mat.sum() == 2
-
-
-def test_integer_confusion_explicit_cap_overrides_local_range():
-    # Even though this pair's own max is 1, an explicit shared cap forces a
-    # wider (and so cross-rollout-consistent) label set.
-    labels, mat = _integer_confusion(np.array([1, 1]), np.array([0, 1]), cap=3)
-    assert labels == ["0", "1", "2", "3+"]
-    assert mat.shape == (4, 4)
-
-
 def test_containment_depths_simple_ramp():
     # one event, edep concentrated in the first bin -> 90%/95% containment
     # depth is the first bin's right edge; a zero-energy event is dropped.
@@ -259,17 +239,25 @@ def test_containment_depths_simple_ramp():
     assert depths.tolist() == [1.0]
 
 
-def test_n_sec_confusion_spec(bundle):
-    spec = get_spec("n_sec_confusion")
+def test_sec_count_per_step_counts_empty_steps(bundle):
+    spec = get_spec("sec_count_per_step")
     r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
-    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
-    assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]]
+    # reference: 3 steps, two of which emit exactly one secondary
+    assert r.payload["reference"][:2] == [1, 2]
+    # rollout: 4 physical steps, one of which emits a single secondary
+    assert r.payload["series"]["rollout"][:2] == [3, 1]
+    assert sum(r.payload["reference"]) == 3
 
 
-def test_n_sec_confusion_shares_one_cap_across_rollouts(two_bundle):
-    spec = get_spec("n_sec_confusion")
-    r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
-    assert list(r.payload["series"]) == ["flow", "wgan"]
-    # both rollouts share the same fixture data here, so their matrices (and
-    # the shared label set) must be identical.
-    assert r.payload["series"]["flow"] == r.payload["series"]["wgan"]
+def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
+    spec = get_spec("sec_count_per_step_by_species")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    cols = r.payload["col_labels"]
+    ref = r.payload["reference"]
+    g = cols.index(pdg_label(22))
+    # two reference steps emit one photon each; the third emits none
+    assert [row[g] for row in ref][:2] == [1, 2]
+    # every other species column is "no such secondary" on all 3 steps
+    for j, _ in enumerate(cols):
+        if j != g:
+            assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0

tests/test_checkpoint_io.py

@@ -14,6 +14,7 @@ from giant import config as gconfig
 from giant.checkpoint_io import (
     CheckpointCompatibilityError,
     InferenceContext,
+    apply_config_overrides,
     conditioning_axes,
     load_for_inference,
     stage_cfg,
@@ -278,3 +279,122 @@ def test_stage_cfg_new_shape_returns_subdict():
 def test_stage_cfg_v02_flat_shape_returns_empty_dict():
     model_cfg = {"hidden_dim": 32, "n_blocks": 4}
     assert stage_cfg(model_cfg, "stage2") == {}
+
+
+# ---------------------------------------------------------------------------
+# config_overrides (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def _router_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
+    return cfg
+
+
+def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.n_sec.sampling": "sample"},
+    )
+    assert ctx.stage2 is not None
+    assert ctx.stage2.n_sec_sampling == "sample"
+    assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}
+
+
+def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
+    )
+    assert ctx.stage1_ddpm_steps == 42
+    assert ctx.stage2_ddpm_steps == 7
+
+
+def test_config_override_other_policy_changes_context_field(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.particle_type.other_policy": "modal"},
+    )
+    assert ctx.other_policy == "modal"
+
+
+def test_config_override_router_temperature_changes_router_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.router.temperature": 1.5},
+    )
+    assert ctx.stage1 is not None
+    assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)
+
+
+def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.config_overrides == {}
+
+
+def test_config_override_unknown_path_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.typo": "sample"},
+        )
+
+
+def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage1_model.hidden_dim": 999},
+        )
+
+
+def test_config_override_bad_value_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.sampling": "maybe"},
+        )
+
+
+def test_apply_config_overrides_no_overrides_returns_same_object():
+    cfg = _model_cfg()
+    assert apply_config_overrides(cfg, None) is cfg
+    assert apply_config_overrides(cfg, {}) is cfg
+
+
+def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
+    legacy_cfg = {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "hidden_dim": 32,
+        "n_blocks": 4,
+        "emb_dim": 8,
+        "dropout": 0.1,
+        "k_max": 5,
+    }
+    merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
+    assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
+    assert merged["stage1_model"]["hidden_dim"] == 32

tests/test_cli_predict.py

@@ -172,3 +172,43 @@ def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+# ---------------------------------------------------------------------------
+# --set (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_set_flag_without_equals_exits_1(tmp_path):
+    checkpoint = tmp_path / "missing.pt"
+
+    result = runner.invoke(
+        app,
+        ["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
+    )
+
+    assert result.exit_code == 1
+    assert "must be 'dotted.path=value'" in result.output
+
+
+def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "predict",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--set",
+            "stage1_model.hidden_dim=999",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_rollout.py

@@ -31,3 +31,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+            "--set",
+            "stage2_model.n_sec.typo=sample",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_train_overrides.py

@@ -20,7 +20,7 @@ def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dic
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["cfg"] = cfg
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     result = runner.invoke(
         cli.app,
@@ -125,7 +125,7 @@ def test_stage2_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage1(mon
 
 
 def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
-    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    monkeypatch.setattr("giant.pipeline.run_train_job", lambda *a, **kw: None)
     result = runner.invoke(
         cli.app,
         ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
@@ -140,7 +140,7 @@ def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_pa
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -161,7 +161,7 @@ def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -178,7 +178,7 @@ def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypat
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.chdir(tmp_path)
 
     result = runner.invoke(cli.app, ["train", "dummy.parquet"])
@@ -192,7 +192,7 @@ def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
         captured["batch_size"] = cfg["train"]["batch_size"]
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
 
     result = runner.invoke(

tests/test_config.py

@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
 def test_n_sec_config_owner_round_trips():
     n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
     assert n_sec.owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
 
 
-def test_n_sec_config_stop_sampling_defaults_to_greedy():
-    assert gconfig.NSecConfig().stop_sampling == "greedy"
+def test_n_sec_config_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().sampling == "greedy"
 
 
-def test_n_sec_config_stop_sampling_round_trips():
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
     n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
-    assert n_sec.stop_sampling == "sample"
-    assert n_sec.to_dict()["stop_sampling"] == "sample"
+    assert n_sec.sampling == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
 
 
 # ---------------------------------------------------------------------------
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
         assert "stop_token" in str(e) and "owner" in str(e)
 
 
-def test_validate_config_bad_stop_sampling_rejected():
-    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+def test_validate_config_bad_n_sec_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_sampling" in str(e)
+        assert "sampling" in str(e)
 
 
 def test_validate_config_default_precision_is_fp32():

tests/test_model_summary.py

@@ -12,6 +12,8 @@ from giant.cli import app
 from giant.materials import MATERIAL_PROPERTIES
 from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
 
+INFERENCE_OVERRIDES = gconfig.INFERENCE_OVERRIDES
+
 runner = CliRunner()
 
 _PDG_VOCAB = 300
@@ -53,6 +55,16 @@ def test_not_build_time_allow_list_has_no_stale_entries():
     assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
 
 
+def test_inference_overrides_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(INFERENCE_OVERRIDES) - in_scope
+    assert not stale, f"INFERENCE_OVERRIDES entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_default_config_overridable_lists_every_allowlisted_path(default_summary):
+    assert set(default_summary.overridable) == set(INFERENCE_OVERRIDES)
+
+
 def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
     assert "stage1_model.router.n_experts" in default_summary.inert
     assert "stage1_model.router.temperature" in default_summary.inert
@@ -135,3 +147,5 @@ def test_cli_default_smoke():
     assert "parameters" in result.output
     assert "trunk" in result.output
     assert "inert under this config" in result.output
+    assert "inference-overridable without retraining" in result.output
+    assert "stage2_model.n_sec.sampling" in result.output

tests/test_render.py

@@ -304,13 +304,15 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "hm1",
             "secondaries",
             "heatmap",
-            "Confusion (single rollout)",
+            "Heatmap (single rollout)",
             "predicted",
             {
                 "series": {"flow": [[1, 0], [0, 1]]},
+                "reference": [[2, 0], [0, 1]],
                 "row_labels": ["0", "1+"],
                 "col_labels": ["0", "1+"],
                 "cbar_label": "count",
+                "log_color": True,
             },
         ),
     ]

tests/test_sample.py

@@ -14,6 +14,7 @@ from giant.model.network import (
     stage2_trunk_sec_dim,
 )
 from giant.sample import (
+    resolve_n_sec,
     sample_flow,
     sample_secondaries,
     sample_secondaries_ar,
@@ -72,6 +73,7 @@ def _stage2_ar(
     mat: int = 2,
     k_max: int = 5,
     history: str = "markov",
+    n_sec_sampling: str = "greedy",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
@@ -89,6 +91,7 @@ def _stage2_ar(
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
 def _stage2_ar_stop_token(
     target: str,
     generator: str,
-    stop_sampling: str = "greedy",
+    n_sec_sampling: str = "greedy",
     emb_dim: int = 6,
     pdg: int = 3,
     mat: int = 2,
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
         particle_type_cfg=ParticleTypeConfig(target=target),
         build_n_sec_head=False,
         build_stop_head=True,
-        stop_sampling=stop_sampling,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
 # ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
     """A stop_head pinned to a large positive logit fires at slot 0 for
     every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
     a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
     generating any token."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, 50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
     assert not sec_valid.any()
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
     """A stop_head pinned to a large negative logit never fires under either
     policy, so every row is capped at k_max (the safety cap, not a modeling
     ceiling)."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, -50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
     stage1_out = torch.randn(3, X_DIM)
     with pytest.raises(AssertionError):
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+
+
+# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
+
+
+def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
+    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
+    returns `bias` (broadcast over the batch) as logits regardless of
+    conditioning — mirrors `_force_stop_head_logit`."""
+    assert decoder.n_sec_head is not None
+    last_linear = decoder.n_sec_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.copy_(bias)
+
+
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
+    """A logit vector overwhelmingly favoring one class gives the same
+    answer under both policies — greedy because it's the argmax, sample
+    because softmax puts ~all mass on it."""
+    B, k_max = 8, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
+    bias = torch.full((k_max + 1,), -50.0)
+    bias[2] = 50.0
+    _force_n_sec_head_bias(decoder, bias)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
+
+
+def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
+    B, k_max = 32, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() == 1
+
+
+def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
+    """Under a flat logit vector, a categorical draw across a large batch
+    should hit more than one class — the whole point of gitea #86: greedy
+    always collapses to one, sample should not."""
+    torch.manual_seed(0)
+    B, k_max = 256, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() > 1