CODING CHALLENGE · N°56

Constrained Decoding (Logit Masking)

Easy AIStructured OutputDecoding

How do you force an LLM to emit only valid JSON, a legal move, or a token your grammar allows? You mask the logits: set every disallowed token to −∞ before picking, so the model can only choose from the permitted set. It’s the backbone of structured-output and grammar-constrained decoding. Solve it in Python or TypeScript, with hidden tests.

The problem

Given the model’s logits (one score per token id) and a list of allowed token ids (the tokens a grammar or schema permits at this step), pick the next token by masking: only the allowed tokens are eligible, and you choose the one with the highest logit. Break ties by the smaller token id. Implement constrained_argmax(logits, allowed): return the chosen token id.

EXAMPLE 1

Input logits = [1.0, 5.0, 2.0, 4.0], allowed = [0, 2, 3]

Output 3

token 1 has the highest logit but is not allowed; among {0,2,3} token 3 (logit 4) wins

EXAMPLE 2

Input logits = [9, 1, 1], allowed = [2, 1]

Output 1

tokens 1 and 2 tie at logit 1 → pick the smaller id, 1

CONSTRAINTS

Only tokens in allowed may be chosen — this is equivalent to setting every other logit to −∞ before taking the argmax.
Among allowed tokens, choose the one with the maximum logit; on a tie, choose the smallest token id.
allowed is non-empty and its ids are valid indices into logits.

SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement constrained_argmax(logits, allowed): scan the allowed ids in increasing order, tracking the best one; update only on a strictly larger logit so ties keep the smaller id. Return the best id.

HINTS — 4 IDEAS

You only ever consider tokens in allowed — the rest are effectively masked to −∞.
Iterate the allowed ids in ascending order and keep the one with the highest logit so far.
Update the best only when you find a strictly larger logit; that way an equal logit does not displace the smaller id already held.
Return the surviving best id — the token constrained decoding would emit.

CPython · WebAssembly

Approach, complexity & discussion — open after you solve

The approach

At each generation step, mask the logits so only tokens the constraint allows can be sampled — set every disallowed token’s logit to −∞ before the softmax, so its probability becomes zero and the model must choose among valid continuations. Advance the constraint’s state (grammar position, remaining set) after each accepted token.

Complexity

O(V) per step to build the mask over the vocabulary (less with a precompiled automaton over allowed tokens).

Common mistakes

Filtering after sampling instead of masking before — that wastes probability mass and can loop retrying.
An over-restrictive mask that leaves no valid token (a dead end).
Not advancing the constraint state per token, so the mask stops matching where you are in the structure.

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.

Finished this one? 0 / 75 Challenges done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

LLM Engineering