Temperature Softmax

Problem

Softmax with temperature scaling.

Scale every logit by 1/temperature, subtract the max for numerical stability, exponentiate, and normalize. Lower temperatures sharpen the distribution and higher temperatures flatten it; assume temperature > 0.

Hint

Dividing the logits by temperature before softmax does the whole job.

Test cases

case 1

in: [[2,4,6],2]

expected: [0.09003057317038046,0.24472847105479764,0.6652409557748218]

case 2

in: [[1,2,3],1]

expected: [0.09003057317038046,0.24472847105479764,0.6652409557748218]

case 3

in: [[1,2,3],0.5]

expected: [0.015876239976466765,0.11731042782619838,0.8668133321973349]

case 4

in: [[0,0],1]

expected: [0.5,0.5]

Starter codePython

import math
def temperature_softmax(logits, temperature):
    # Your code here
    pass

4 test cases · run your code against them in the editor