Log-Softmax Jacobian Matrix

Problem

For y = log softmax(z), the Jacobian is I - 1 s^T, where s = softmax(z): diagonal entries 1 - s_i and off-diagonal entries -s_j.

Given the logits, return the Jacobian as nested rows.

Hint

Each row i subtracts the same softmax vector from the identity row.

Test cases

case 1

in: [[0,0]]

expected: [[0.5,-0.5],[-0.5,0.5]]

case 2

in: [[1,2,3]]

expected: [[0.9099694268296196,-0.24472847105479764,-0.6652409557748218],[-0.09003057317038046,0.7552715289452023,-0.6652409557748218],[-0.09003057317038046,-0.2447284710…

case 3

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

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

Starter codePython

def log_softmax_jacobian(logits):
    # Your code here
    pass

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