Add sampled n_sec under n_sec.mode = 'head' (gitea #86)

Taking argmax over the n_sec classifier logits collapses secondary
multiplicity onto its conditional mode at fixed pre-step conditioning,
under-dispersing n_sec in rollouts and biasing low wherever the true
conditional count distribution is right-skewed (typical for
multiplicity).

Generalizes stage2_model.n_sec.stop_sampling (previously stop_token-only)
into stage2_model.n_sec.sampling, covering both "head" (greedy: argmax;
sample: categorical draw via torch.multinomial) and "stop_token" (unchanged:
greedy threshold / Bernoulli draw) modes. stop_sampling is kept as a
deprecated alias in NSecConfig.from_dict and migrate_config, since it
appears in existing checkpoints' model_config. Default stays "greedy" so
existing runs/checkpoints are unaffected.

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