Derivative of Softmax

Calculus & Optimization DS practice problem on Onlearn.

Difficulty: hard.

Topics: Understanding the Jacobian of the Softmax Function, Softmax Function, Kronecker Delta, Outer Product, Diagonal Matrix, Vectorized Matrix Subtraction, Multivariable Calculus, Linear Algebra, Probability Theory, Neural Network Optimization, Numerical Stability, Partial Derivatives, Jacobian Matrices, Gradient Descent, Log-Sum-Exp Trick, Chain Rule.

Implement a function that computes the Jacobian matrix of the softmax function for a given input vector z. The Jacobian J is an N x N matrix where J {ij} = d(softmax(z) i) / d(z j). Your output should return the full matrix.