Softmax Linear Jacobian

Problem

For y = softmax(W x) with a k x n weight matrix W, the Jacobian with respect to x is (diag(s) - s s^T) W, where s is the softmax of Wx.

Given W as nested rows and x, return the k x n Jacobian as nested rows.

Hint

Build diag(s) - s s^T first, then multiply by W.

case 1

in: [[[1,0],[0,1]],[0,0]]

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

case 2

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

expected: [[0]]

case 3

in: [[[1,0],[0,1]],[1,-1]]

expected: [[0.10499358540350662,-0.10499358540350649],[-0.10499358540350649,0.1049935854035065]]

Starter codePython

def softmax_linear_jacobian(w, x):
    # Your code here
    pass

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