Start from brute force
The obvious approach follows the definition literally: copy logits, set every id not in allowed to −∞, then take the argmax over the whole vector, keeping the first maximum so ties go to the smaller id. It is correct because a −∞ logit loses every comparison. It costs O(V) time and O(V) space for the masked copy, where V is the vocabulary size, plus O(a) to build a set of the a allowed ids. The waste: every decoding step touches all V logits, often tens of thousands, even when the grammar allows only a handful of tokens, such as a digit or a closing brace.
The approach
Masking to −∞ and then taking the argmax is the same as taking the argmax over the allowed ids alone, so the mask never needs to exist. The reference walks sorted(allowed) in increasing id order and keeps best, replacing it only when logits[idx] > logits[best]. The strict comparison is the tie-break: an equal logit later in the scan always has a larger id, so it cannot displace the smaller id already held. Starting from best = None rather than a score of 0 matters too, because every allowed logit can be negative.
The invariant after each step: best is the smallest id carrying the maximum logit among the ids seen so far. The sort costs O(a log a); a single unsorted pass that also accepts an equal logit with a smaller id gives the same answer in O(a).
Walk through an example
With logits = [2.0, 7.0, 3.5, 3.5, 6.0, 1.0] and allowed = [5, 3, 2, 0], the reference scans sorted(allowed) = [0, 2, 3, 5], never looks at the two highest logits (ids 1 and 4), and resolves the 3.5 tie between ids 2 and 3 to 2.
| step | idx | logits[idx] | comparison | best |
|---|
| 1 | 0 | 2.0 | best is None, so take it | 0 |
| 2 | 2 | 3.5 | 3.5 > 2.0 | 2 |
| 3 | 3 | 3.5 | 3.5 > 3.5 is false, so the tie keeps 2 | 2 |
| 4 | 5 | 1.0 | 1.0 > 3.5 is false | 2 |
Complexity
The reference costs O(a log a) time for the sort and O(a) space for the sorted copy, where a = len(allowed); a one-pass tie-break makes the time O(a). Masking the whole vocabulary costs O(V) time and O(V) space at every decoding step, however few tokens are allowed, because it builds and scans a full-size copy.
Edge cases to test
- Every allowed logit is negative:
constrained_argmax([-1.0, 5.0, -3.0], [0, 2]) returns 0. A solution that seeds the best score with 0 never updates and returns its sentinel. - A tie:
constrained_argmax([9, 1, 1], [2, 1]) returns 1. Scanning allowed in the given order with > returns 2, and so does a sorted scan with >=. - One allowed id:
constrained_argmax([5, 9, 2], [0]) returns 0, even though token 1 holds the top logit. The grammar has already decided. - Duplicate ids in
allowed are harmless: constrained_argmax([1, 2, 3], [2, 2, 0]) returns 2. - A NaN logit compares false against everything, so it is never picked over a number but never displaced once held:
constrained_argmax([float("nan"), 1.0], [0, 1]) returns 0. Reject NaN logits before decoding.
Common mistakes
- Taking the argmax over all logits and only then checking whether it is allowed. In the first example the top token, 1, is not allowed, and there is no second choice to fall back to.
- Scanning
allowed in the order given instead of ascending id order, so a tie goes to whichever id is listed first. - Using
>= instead of >, which hands every tie to the larger id. - Seeding the best score with 0 instead of the first allowed logit, which breaks when every allowed logit is negative.
Follow-ups an interviewer asks next
- How does this extend to generating a whole JSON object?
- Call it once per decoding step with a fresh
allowed set. A grammar or regex compiled to a finite-state machine tracks the position in the structure; after each emitted token the state advances and yields the next step’s allowed ids. This function is the masking step. The hard part is computing those sets quickly when tokens span several characters. - What changes if you sample instead of taking the argmax?
- Set disallowed logits to
−∞, apply softmax (with temperature) to the rest, and sample. Since exp(−∞) = 0, masked tokens get probability exactly 0 and the allowed ones renormalize among themselves: for the first example, about 0.844 for token 3, 0.114 for token 2 and 0.042 for token 0. Top-k or top-p then filters the masked distribution. - What if the grammar leaves no allowed token?
- That is a dead end: the constraint and the text so far are inconsistent, often because the grammar has a bug or a token boundary split a required string. This function assumes
allowed is non-empty; given an empty list the reference returns None (TypeScript: -1). A production decoder should detect the empty set and fail loudly or backtrack.
Where this shows up
Constrained decoding is how you guarantee structured output — valid JSON, a value from an enum, a well-formed SQL query — instead of hoping the model complies. It underpins grammar-constrained generation and the reliable, schema-valid arguments that make tool-calling trustworthy.
Reference solution
Reference solution — Python
def constrained_argmax(logits, allowed):
best = None
for idx in sorted(allowed): # ascending -> ties keep smaller id
if best is None or logits[idx] > logits[best]:
best = idx # strictly larger logit wins
return best
print(constrained_argmax([1.0, 5.0, 2.0, 4.0], [0, 2, 3])) # -> 3
print(constrained_argmax([9, 1, 1], [2, 1])) # -> 1
Reference solution — TypeScript
function constrainedArgmax(logits: number[], allowed: number[]): number {
let best = -1;
for (const idx of [...allowed].sort((a, b) => a - b)) { // ascending -> ties smaller id
if (best === -1 || logits[idx] > logits[best]) best = idx; // strictly larger wins
}
return best;
}
console.log(constrainedArgmax([1.0, 5.0, 2.0, 4.0], [0, 2, 3])); // -> 3
console.log(constrainedArgmax([9, 1, 1], [2, 1])); // -> 1