Validate top_k and top_p before the greedy branch in sample_logits

tests/unit/utilities/test_logits_utils.py

@@ -112,9 +112,12 @@ def test_top_k_restricts_to_dominant_token(self):
         outs = [sample_logits(logits, top_k=1).item() for _ in range(20)]
         assert set(outs) == {1}
 
-    def test_top_k_rejects_non_positive(self):
-        with pytest.raises(AssertionError):
-            sample_logits(torch.randn(1, 4), top_k=0)
+    @pytest.mark.parametrize("temperature", [1.0, 0.0])
+    @pytest.mark.parametrize("top_k", [0, -5])
+    def test_top_k_rejects_non_positive(self, temperature, top_k):
+        # Greedy sampling skips the filtering branch, so the check has to come before it.
+        with pytest.raises(AssertionError, match="top_k has to be greater than 0"):
+            sample_logits(torch.randn(1, 4), top_k=top_k, temperature=temperature)
 
 
 class TestSampleLogitsTemperature:
@@ -172,11 +175,12 @@ def test_top_p_keeps_dominant_token(self):
         outs = [sample_logits(logits, top_p=0.5).item() for _ in range(20)]
         assert set(outs) == {1}
 
-    def test_top_p_rejects_out_of_range(self):
-        with pytest.raises(AssertionError):
-            sample_logits(torch.randn(1, 4), top_p=0.0)
-        with pytest.raises(AssertionError):
-            sample_logits(torch.randn(1, 4), top_p=1.5)
+    @pytest.mark.parametrize("temperature", [1.0, 0.0])
+    @pytest.mark.parametrize("top_p", [0.0, -0.5, 1.5])
+    def test_top_p_rejects_out_of_range(self, temperature, top_p):
+        # Same gap as top_k: a greedy call used to accept any top_p without complaint.
+        with pytest.raises(AssertionError, match=r"top_p has to be in \(0, 1\]"):
+            sample_logits(torch.randn(1, 4), top_p=top_p, temperature=temperature)
 
 
 class TestSampleLogitsFreqPenalty:

transformer_lens/utilities/logits_utils.py

@@ -107,6 +107,11 @@ def sample_logits(
     if repetition_penalty <= 0:
         # Zero divides logits by zero (NaN) and a negative value flips their signs.
         raise ValueError(f"repetition_penalty must be positive, got {repetition_penalty}.")
+    if top_k is not None:
+        # Checked before the greedy branch, so an invalid filter is never silently ignored.
+        assert top_k > 0, "top_k has to be greater than 0"
+    if top_p is not None:
+        assert 1.0 >= top_p > 0.0, "top_p has to be in (0, 1]"
     if temperature == 0.0:
         # Greedy sampling - still apply repetition penalty before argmax
         if repetition_penalty != 1.0 and tokens is not None:
@@ -133,7 +138,6 @@ def sample_logits(
                     tokens[batch_index], minlength=final_logits.shape[-1]
                 )
         if top_k is not None:
-            assert top_k > 0, "top_k has to be greater than 0"
             # Clamp top_k to the vocab size so a large value does not raise
             # "selected index k out of range" (matches HuggingFace's
             # TopKLogitsWarper, which does top_k = min(top_k, logits.size(-1))).
@@ -142,7 +146,6 @@ def sample_logits(
             indices_to_remove = final_logits < top_logits[..., -1].unsqueeze(-1)
             final_logits = final_logits.masked_fill(indices_to_remove, -float("inf"))
         elif top_p is not None:
-            assert 1.0 >= top_p > 0.0, "top_p has to be in (0, 1]"
             sorted_logits, sorted_indices = torch.sort(final_logits, descending=True)
             cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
             # We round up - we want prob >= top_p not <top_p