Softmax With Temperature

Problem

A language model scores every candidate next token with a raw number called a logit, and sampling turns those scores into probabilities with a softmax: divide each logit by the temperature T, exponentiate, and normalise so the results sum to 1. Write softmax(logits, temperature) returning the probabilities rounded to 3 decimals. It must not overflow on large logits such as 1000, and a temperature of 0 means greedy decoding, so all of the probability goes to the highest logit, the first one on a tie.

Examples

Input:  logits = [2.0, 1.0, 0.1], temperature = 1.0
Output: [0.659, 0.242, 0.099]

Input:  logits = [2.0, 1.0, 0.1], temperature = 0.5
Output: [0.864, 0.117, 0.019]
Why:    a low temperature stretches the gaps, so the favourite takes more

Input:  logits = [1000.0, 999.0], temperature = 1.0
Output: [0.731, 0.269]
Why:    edge case, exp(1000) overflows, but only the differences between logits matter

Hints

0 / 3

1. Hint 1reveal
2. Hint 2locked
3. Hint 3locked

import math

def softmax(logits, temperature=1.0):
    if temperature == 0:                   # greedy: all the mass on the top logit
        best = logits.index(max(logits))
        return [1.0 if i == best else 0.0 for i in range(len(logits))]
    scaled = [x / temperature for x in logits]
    top = max(scaled)                      # shift by the max so exp never overflows
    exps = [math.exp(x - top) for x in scaled]
    total = sum(exps)
    return [round(e / total, 3) for e in exps]

print(softmax([2.0, 1.0, 0.1]))                  # -> [0.659, 0.242, 0.099]
print(softmax([2.0, 1.0, 0.1], 0.5))             # -> [0.864, 0.117, 0.019]
print(softmax([2.0, 1.0, 0.1], 5.0))             # -> [0.4, 0.327, 0.273]
print(softmax([1000.0, 999.0]))                  # -> [0.731, 0.269]
print(softmax([3.0, 5.0, 5.0], 0))               # -> [0.0, 1.0, 0.0]