Temperature Softmax

Problem

Compute softmax with temperature: softmax(logits / temperature), computed stably by subtracting the maximum scaled logit before exponentiating.

The signature is temperature_softmax(logits, temperature). Assume temperature > 0. The returned probabilities sum to 1.0.

Hint

Divide the logits by the temperature first, then run stable softmax.

case 1

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

expected: [0.0900305732,0.2447284711,0.6652409558]

case 2

in: [[1,1],0.5]

expected: [0.5,0.5]

case 3

in: [[0,0,0],100]

expected: [0.3333333333,0.3333333333,0.3333333333]

case 4

in: [[2,0],2]

expected: [0.7310585786,0.2689414214]

case 5

in: [[1,2],0.1]

expected: [0.0000453979,0.9999546021]

Starter codePython

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

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