Implement a batched top-p filter that keeps tied tokens together and handles temperature 0

mediumFree drillimplement top_p_filter

Before you code

Change the temperature while keeping the logits and top-p threshold fixed. Predict whether the same tokens will remain eligible.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

Write the probability filter a sampler applies to one batch of logits. Each row has its own temperature and top-p.

top_p_filter(logits, temperature, top_p) -> numpy array [B, V] of probabilities
  logits       [B, V] floats; a disallowed token is -inf
  temperature  [B]    0 means greedy
  top_p        [B]    in (0, 1]

The rules, per row:

Temperature 0 returns a one-hot row on the argmax, and the lowest token id wins a tie.
Otherwise divide the logits by the temperature, softmax, and find the smallest set of highest-probability tokens whose mass reaches p. Every token whose probability equals the probability of the last token in that set is kept too, so tied tokens are kept or dropped together.
p of 1 keeps every token with nonzero probability.
Renormalize what is kept so each row sums to 1.

The visible tests cover a tie at the cutoff and the greedy branch. The full suite adds a tie at the argmax, p of 1 over a 128,256-token vocabulary given as float32, a tiny p, the order of temperature and top-p, mixed rows in one batch and masked tokens.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Free drill: the visible tests run for everyone, the edge cases and the solution for any signed-in reader.

The visible tests

import numpy as np

LOGITS = np.array([[2.0, 1.0, 1.0, 0.5, -1.0]])


def test_a_tie_at_the_cutoff_keeps_both_tokens(ns):
    out = np.asarray(ns["top_p_filter"](LOGITS, np.array([1.0]), np.array([0.65])))
    kept = sorted(np.flatnonzero(out[0] > 0).tolist())
    assert kept == [0, 1, 2], f"kept {kept}; ids 1 and 2 have equal probability, so both stay"
    assert abs(out[0].sum() - 1) < 1e-6, "the kept probabilities must sum to 1"
    assert abs(out[0, 0] - 0.5761168847658291) < 1e-6, f"token 0 should be 0.576 after renormalizing, got {out[0, 0]:.4f}"


def test_temperature_zero_is_greedy(ns):
    out = np.asarray(ns["top_p_filter"](LOGITS, np.array([0.0]), np.array([0.9])))
    assert out[0].tolist() == [1.0, 0.0, 0.0, 0.0, 0.0], f"expected a one-hot row on token 0, got {out[0].round(3).tolist()}"

The reasoning behind this component: the companion page.