From f515fe828da541cfaf5875ee98f48e3ebb4cfff7 Mon Sep 17 00:00:00 2001 From: dvir arad Date: Thu, 23 Jul 2026 17:37:50 +0300 Subject: [PATCH] fix(inference): handle None top_k/top_p in sample_from_logits `sample_from_logits` declares `top_k=None` and `top_p=None` defaults, but its guard evaluated `top_k > 0 or top_p < 1.0` whenever either value was not None. Passing only one filter (e.g. `sample_from_logits(logits, top_p=0.9)`) therefore raised `TypeError: '>' not supported between instances of 'NoneType' and 'int'`, and a None value could also be passed down to `top_k_top_p_filtering`, which performs the same comparison. Normalize the unset parameter to its neutral default (top_k=0, top_p=1.0) before the comparison. Behavior is unchanged for all existing call sites, which always pass explicit numeric values. --- model/kronos.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/model/kronos.py b/model/kronos.py index ce4494ee0..eef8db24f 100644 --- a/model/kronos.py +++ b/model/kronos.py @@ -372,9 +372,15 @@ def top_k_top_p_filtering( def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True): logits = logits / temperature - if top_k is not None or top_p is not None: - if top_k > 0 or top_p < 1.0: - logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) + # Normalize optional filtering parameters. Leaving one of them unset (None) + # while passing the other (e.g. sample_from_logits(logits, top_p=0.9)) would + # otherwise raise `TypeError: '>' not supported between 'NoneType' and 'int'` + # in the comparison below, and could also pass None down to + # top_k_top_p_filtering. + top_k = 0 if top_k is None else top_k + top_p = 1.0 if top_p is None else top_p + if top_k > 0 or top_p < 1.0: + logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) probs = F.softmax(logits, dim=-1)