MediumCalculusca-287

Temperature Softmax Jacobian Diagonal

Problem

For temperature-scaled softmax s = softmax(logits / T), the Jacobian with respect to logits has diagonal entries (1/T) * s_i * (1 - s_i).

Given logits and temperature, return this diagonal list. Subtract the max logit before exponentiating.

Hint

Chain rule with u = logits / T brings a 1/T factor into every entry.

case 1

in: [[0,0],1]

expected: [0.25,0.25]

case 2

in: [[0,0],2]

expected: [0.125,0.125]

case 3

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

expected: [0.031248369961352808,0.20709738269886538,0.23089595864457524]

Starter codePython

def temperature_softmax_diag(logits, temperature):
    # Your code here
    pass

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