Stretch: temperature + top-k sampling

The problem, from first principles

[07-greedy-decoding]'s argmax is entirely deterministic: the same input always produces the same output, and it always produces the SINGLE most likely continuation, never anything else, even when several plausible continuations exist with only slightly lower probability. For creative or varied generation (chat responses, story writing, anything where always producing the identical "safest" answer feels repetitive or dull), what's actually wanted is RANDOM sampling from the model's own predicted probability distribution, weighted by that distribution, not a hard argmax.

Two knobs give control over exactly how random: temperature RESHAPES the whole distribution before sampling (dividing logits by temperature, exactly [02-modern-transformer-architecture/10-logit-scaling]'s scaling operation, but user-chosen here rather than fixed by d_model, low temperature sharpens toward argmax-like behavior, high temperature flattens toward uniform randomness), and top-k FILTERS the candidate pool down to only the k highest-scoring tokens before sampling at all, preventing the model from ever picking something from deep in its low-probability tail, however unlikely, purely by chance.

From theory to code

Implement scale_and_filter_logits(logits_row, temperature, top_k) (divide by temperature, then mask everything outside the top k to -inf), sample_next_token(logits_row, temperature, top_k, rng) (scale_and_filter_logits followed by [04-seq-modeling/04-attention/03-softmax-last-axis]'s softmax and a random draw), and sample_decode, [07-greedy-decoding]'s autoregressive loop with sample_next_token replacing argmax.

Constraints

• temperature divides the logits BEFORE any top-k filtering; top_k filtering happens on the ALREADY-scaled logits.
• top_k=None means no filtering at all (sample from the full distribution); top_k=1 reduces to [07-greedy-decoding]'s deterministic argmax behavior, since only the single highest-scoring token remains eligible.
• Masked-out (below top-k) logits become -inf, giving them exactly 0 probability after softmax, never merely a SMALL probability.
• sample_next_token draws using rng.choice (an explicitly-passed random state, for reproducibility across calls), never NumPy's unseeded global randomness.

Hints

scaled = logits_row / temperature
if top_k is not None and top_k < len(scaled):
    threshold = np.sort(scaled)[-top_k]
    scaled = np.where(scaled >= threshold, scaled, -np.inf)
return scaled

probs = softmax_last_axis(scale_and_filter_logits(logits_row, temperature, top_k))
return int(rng.choice(len(probs), p=probs))