CODING CHALLENGE · N°52

Softmax with Temperature

Easy AISamplingMath

The dial that controls how "creative" a language model is: temperature scales the logits before the softmax. Low temperature sharpens the distribution toward the top token (greedy, deterministic); high temperature flattens it toward uniform (random, diverse). Implement it — numerically stable. Solve it in Python or TypeScript, with hidden tests.

The problem

Implement softmax_temp(logits, temperature): divide each logit by temperature, then apply the softmax to turn the scaled logits into a probability distribution (positive values summing to 1). Do it in a numerically stable way by subtracting the max before exponentiating. Lower temperature → sharper (more peaked) distribution; higher temperature → flatter.

EXAMPLE 1
Input logits = [5, 5, 5, 5], temperature = 2.0
Output [0.25, 0.25, 0.25, 0.25]
equal logits give a uniform distribution at any temperature
EXAMPLE 2
Input logits = [2, 1, 0], temperature = 0.5
Output ≈ [0.867, 0.117, 0.016]
low temperature concentrates mass on the top logit
CONSTRAINTS
  • Return a list of probabilities that are all ≥ 0 and sum to 1.
  • Subtract max(scaled_logits) before exp to avoid overflow — this does not change the result, only its stability.
  • Temperature < 1 sharpens the distribution (toward argmax); temperature > 1 flattens it (toward uniform).
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement softmax_temp(logits, temperature): scale by dividing by the temperature, subtract the max for stability, exponentiate, and normalise by the sum.

HINTS — 4 IDEAS
  1. First scale: z[i] = logits[i] / temperature.
  2. Subtract the maximum of z from every element before exponentiating — this prevents overflow and leaves the probabilities unchanged.
  3. Exponentiate each shifted value, then divide by the sum of all of them so the result sums to 1.
  4. Sanity check: at high temperature the outputs approach uniform; at low temperature the largest logit dominates.
CPython · WebAssembly

Explore the topic

See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.