fix(sample_utils): correct top_k error message to the exclusive (0, vocab_size) bound (ml-explore#1377)

devYRPauli · web-flow · commit df48987708fc · 2026-06-09T11:46:34.000+02:00
diff --git a/mlx_lm/sample_utils.py b/mlx_lm/sample_utils.py
@@ -141,7 +141,7 @@ def apply_top_k(
     vocab_size = logprobs.shape[-1]
     if not isinstance(top_k, int) or not (0 < top_k < vocab_size):
         raise ValueError(
-            f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
+            f"`top_k` has to be an integer in the (0, {vocab_size}) interval,"
             f" but is {top_k}."
         )
     mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]

Original file line number	Diff line number	Diff line change
`@@ -141,7 +141,7 @@ def apply_top_k(`
`141`	`141`	`vocab_size = logprobs.shape[-1]`
`142`	`142`	`if not isinstance(top_k, int) or not (0 < top_k < vocab_size):`
`143`	`143`	`raise ValueError(`
`144`		- f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
	`144`	+ f"`top_k` has to be an integer in the (0, {vocab_size}) interval,"
`145`	`145`	`f" but is {top_k}."`
`146`	`146`	`)`
`147`	`147`	`mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]`