Derivative of the log softmax function

I am trying to find the derivative of the log softmax function :

$$LS(z) = log(\frac {e^{z-c}} {\sum_{i=0}^n e^{z_i-c}}) = z - c - log({\sum_{i=0}^n e^{z_i-c}})$$

(c = max(z) )

with respect to the input vector z. However it seems I have made a mistake somewhere. Here is what I have attempted out so far:

The partial derivative of logsoftmax with respect to an element of z

$${\partial {LS}\over\partial{z_i}}= 1 - {\partial{ log({\sum_{i=0}^n e^{z_i-c}})} \over \partial {z_i} }$$

$$= 1 - {1\over{{\sum_{i=0}^n e^{z_i-c}}}} * {\partial{ {\sum_{i=0}^n e^{z_i-c}}}\over \partial z_i}$$

$$= 1 - {e^{z-c}\over {\sum_{i=0}^n e^{z_i-c}}}$$

$$= 1 - Softmax(z)_i$$

Now, Since I am trying to use this for autodifferentiation, I can then substitute the Softmax with logsoftmax for efficiency since I already compute it in the forward pass:

$$= 1 - e^{LS(z)_i}$$

So in vectorized form, we then get:

$${\partial {LS}\over\partial{z}} = \vec 1 - \vec {e^{LS(z)}}$$

In numpy code:

def logsoftmax(z):
c = np.max(z, axis=1)
lexpsum = c + np.log(np.exp(z - c.reshape((-1,1))).sum(axis=1))
out = z - lexpsum.reshape((-1,1))

# out is the logsoftmax we computed and out.grad is the incoming gradient which will be
# multiplied to this gradient as per the chain rule
return out.grad * (1 - np.exp(out))

# z = np.array([[1., 2., 3.]])
# My output:
# logsoftmax: [[-2.4076061 , -1.4076061 , -0.40760612]]
# logsoftmax_backward: [[0.90996945, 0.75527155, 0.33475912]]
# Gradient using pytorch: [[ 0.7299,  0.2658, -0.9957]]


It seems I am missing a normalizing constant somewhere, or I have done something completely wrong, but I can't figure it out.

$$\def\K{{\large\ell}} \def\o{{\tt1}}\def\p{\partial} \def\L{\left}\def\R{\right} \def\LR#1{\L(#1\R)} \def\BR#1{\Big(#1\Big)} \def\vec#1{\operatorname{vec}\LR{#1}} \def\diag#1{\operatorname{diag}\LR{#1}} \def\Diag#1{\operatorname{Diag}\LR{#1}} \def\trace#1{\operatorname{Tr}\LR{#1}} \def\qiq{\quad\implies\quad} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\c#1{\color{red}{#1}}$$Denote the all-ones vector by $$\o$$ and define an offset vector $$y\;=\;z-c \quad\;\qiq dy = dz$$ Then calculate the gradient of the element-wise exponential function \eqalign{ p &= \exp(y) &\qiq P = \Diag{p} \\ dp &= P\,dy = P\,dz \\ \grad{p}{z} &= P \\ } Next, calculate the gradient of the element-wise softmax function \eqalign{ s &= \frac{p}{\o^Tp} &\qiq S = \Diag{s} = \frac {P}{\o^Tp} \\ ds &= \frac{dp}{{\o^Tp}} - \frac{p\LR{\o^Tdp}}{\LR{\o^Tp}^2} \\ &= \frac{P\,dz}{{\o^Tp}} - \frac{s\LR{\o^TP\,dz}}{{\o^Tp}} \\ &= \BR{S-ss^T}\,dz \\ \grad{s}{z} &= \BR{S-ss^T} \\ } Finally, tackle the element-wise logarithm of the softmax function \eqalign{ \K &= \log(s) \\ d\K &= S^{-1}\,ds \\ &= S^{-1}\BR{S-ss^T}\,dz \\ &= \BR{I-\o s^T}\,dz \\ \grad{\K}{z} &= \BR{I-\o s^T} \\ }