Fix XTC sampling threshold to be per-row for batched logits

mlx_lm/sample_utils.py

@@ -263,7 +263,9 @@ def apply_xtc(
         )
 
     probs = mx.softmax(logits, -1)
-    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
+    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min(
+        axis=-1, keepdims=True
+    )
     if xtc_special_tokens:
         mask[..., xtc_special_tokens] = False
 

tests/test_sample_utils.py

@@ -116,6 +116,13 @@ def test_apply_xtc(self):
         new_probs = mx.softmax(apply_xtc(mx.log(probs), 0, 0.1, [0]), -1)
         self.assertTrue(mx.allclose(new_probs, probs))
 
+        # Test that the threshold is computed per row in a batch
+        probs = mx.array([[0.6, 0.25, 0.1, 0.05], [0.9, 0.05, 0.03, 0.02]])
+        batched = mx.softmax(apply_xtc(mx.log(probs), 1, 0.2, []), -1)
+        for i in range(probs.shape[0]):
+            alone = mx.softmax(apply_xtc(mx.log(probs[i : i + 1]), 1, 0.2, []), -1)
+            self.assertTrue(mx.allclose(batched[i : i + 1], alone))
+
     def test_presence_penalty(self):
         from mlx_lm.sample_utils import make_presence_penalty