Vibe Engines Learn
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
  1. You only ever consider tokens in allowed — the rest are effectively masked to −∞.
  2. Iterate the allowed ids in ascending order and keep the one with the highest logit so far.
  3. Update the best only when you find a strictly larger logit; that way an equal logit does not displace the smaller id already held.
  4. Return the surviving best id — the token constrained decoding would emit.
CPython · WebAssembly

Explore the topic

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

LLM Engineering