Softmax Jacobian via Chain Rule

Problem

For the softmax vector s, the full Jacobian is J[i][j] = s_i * (delta_ij - s_j), where delta_ij is 1 when i == j and 0 otherwise.

Given the logits, return the n x n Jacobian as a nested list. Use max subtraction for numerical stability.

Hint

Each row sums to zero because the softmax outputs always sum to 1.

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: [[1,1,1]]

expected: [[0.22222222222222224,-0.1111111111111111,-0.1111111111111111],[-0.1111111111111111,0.22222222222222224,-0.1111111111111111],[-0.1111111111111111,-0.11111111111…

case 4

in: [[2,-1]]

expected: [[0.045176659730912,-0.045176659730912144],[-0.045176659730912144,0.045176659730912144]]

Starter codePython

def softmax_jacobian(logits):
    # Your code here
    pass

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