Top-k Sampling

Keep the k highest-scoring tokens, then draw one token according to their relative probabilities. Filtering chooses which tokens are eligible; sampling chooses one of those tokens.

Your task

Implement top_k_sample(logits, k, seed).

• logits is a one-dimensional floating-point tensor of finite scores, with one score per vocabulary token. A token's ID is its index.
• k is an integer with 1 <= k <= len(logits).
• seed initializes the random generator for this single draw.
• Return one Python integer token ID. Leave logits unchanged.

Rank tokens by descending score, breaking ties by lower token ID. Retain exactly the first k tokens in that ordering. A stable descending sort gives this tie rule; an unspecified top-k tie order does not.

Keep the original vocabulary positions and replace excluded scores with negative infinity. Softmax then assigns excluded tokens exactly zero probability and renormalizes the retained scores. Do not add an epsilon to excluded probabilities, which would make forbidden tokens eligible again.

For logits [3, 2, 0] and k=2, only IDs 0 and 1 are eligible. Their probabilities are approximately [0.731, 0.269]; ID 2 has probability zero. The largest score is more likely, but a draw can still select ID 1.

Reproduce the framework's draw

The fixtures check a seeded draw, not just whether the returned token belongs to the retained set. Use the full vocabulary-sized distribution:

In PyTorch

Select indices with torch.argsort(logits, descending=True, stable=True)[:k], mask the other logits, apply softmax and call torch.multinomial with num_samples=1 and a fresh torch.Generator().manual_seed(seed).

In JAX

Select indices with jnp.argsort(-logits, stable=True)[:k], mask the other logits and call jax.random.categorical(jax.random.PRNGKey(seed), filtered). Categorical sampling accepts logits directly, so no softmax/log round trip is needed.

Convert the sampled index to a Python int. These seeded fixtures use float32 inputs and a CPU generator for PyTorch. JAX fixtures use the runtime's jax_threefry_partitionable=True setting. Different frameworks need not draw the same token from the same seed; each has its own expected outputs. Reordering or shrinking the sampling distribution can also change seeded results.

With k=1, the retained token is always selected. With k equal to vocabulary size, all tokens are eligible. Limiting the candidate set does not guarantee coherent text. This function performs one draw, not a generation loop.

The diagram illustrates how scores become probabilities. This exercise has no temperature parameter and uses the supplied scores without rescaling them.