Softmax Jacobian Diagonal

Problem

For softmax probabilities s = softmax(logits), the diagonal of the Jacobian contains the entries s_i * (1 - s_i).

Given the logits, return this diagonal as a list. Subtract the maximum logit before exponentiating for numerical stability.

Hint

Each diagonal entry is the variance of a Bernoulli with success probability s_i.

case 1

in: [[0,0]]

expected: [0.25,0.25]

case 2

in: [[0,1]]

expected: [0.19661193324148185,0.19661193324148185]

case 3

in: [[1,1,1]]

expected: [0.22222222222222224,0.22222222222222224,0.22222222222222224]

case 4

in: [[2,-1]]

expected: [0.045176659730912,0.045176659730912144]

Starter codePython

def softmax_jacobian_diagonal(logits):
    # Your code here
    pass

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