Skip to content

← JAX Stochasticity step 9 of 25

Medium Primitives

Top-k Logit Masking

Jump to the task

Why this matters

A top-k mask limits which token scores can contribute to a later sampling step. This exercise finds the k-th largest score and keeps every score at least that large, including ties.

Masking is separate from softmax and sampling. The output here is still a vector of logits in its original order.

Worked mini-example

import jax
import jax.numpy as jnp

logits = jnp.array([1.0, 5.0, 3.0, 2.0, 4.0])
values, indices = jax.lax.top_k(logits, 2)
threshold = values[-1]  # 4
masked = jnp.where(logits >= threshold, logits, -1e9)
# [-1e9, 5, -1e9, -1e9, 4]

Use the selected values to find the cutoff, then apply that cutoff to the original vector. Sorting the returned logits would lose their correspondence with token indices.

Common pitfalls

  • This task keeps all ties at the cutoff. With scores [5, 4, 4, 1] and k = 2, it keeps three entries. That is different from selecting exactly two indices.
  • The finite sentinel -1e9 is part of this exercise's output contract. It is not mathematically equivalent to negative infinity. If retained scores are also extremely negative, the sentinel might not suppress masked probabilities.
  • Using negative infinity can give exactly zero probability after softmax when a finite score remains. A finite sentinel may only approximate that behavior, depending on scale and dtype.
  • Convert the concrete count with int(k). The output size of lax.top_k depends on k, so an outer JIT must make k static; casting a dynamic tracer does not make it concrete.
  • Use >= at the cutoff so the k-th value and its ties survive.

Problem

Implement top_k_mask(logits, k). Find the k-th largest value, retain every logit greater than or equal to it, and replace each smaller logit with exactly -1e9.

logits is a nonempty one-dimensional array of finite float32 values with length K. The count is an integer-valued Python float satisfying 1 <= k <= K; convert it to an integer.

Return the same shape and preserve the original order. For [3, 1, 2] and k = 1.0, return [3, -1e9, -1e9]. When all scores tie, keep them all. When k = K, keep every score.

The runner calls this function directly. To JIT it yourself, use jax.jit(top_k_mask, static_argnames=("k",)).