MediumLinear Algebrala-205

Softmax Jacobian

Problem

Return the Jacobian of the softmax function at vector v:

J = diag(s) - s*s^T

where s = softmax(v), so J[i][j] = s[i] * (delta_ij - s[j]) with delta the Kronecker delta. Use a stable softmax.

Hint

The softmax Jacobian is symmetric and every row sums to 0.

Test cases

case 1

in: [[0,0]]

expected: [[0.25,-0.25],[-0.25,0.25]]

case 2

in: [[0,1]]

expected: [[0.19661193324148185,-0.19661193324148185],[-0.19661193324148185,0.19661193324148185]]

case 3

in: [[1000,1000]]

expected: [[0.25,-0.25],[-0.25,0.25]]

case 4

in: [[1,2,3]]

expected: [[0.08192506906499324,-0.022033044520174298,-0.05989202454481893],[-0.022033044520174298,0.1848364465099787,-0.1628034019898044],[-0.05989202454481893,-0.162803…

Starter codePython

def softmax_jacobian(v):
    # Your code here
    pass

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