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))

def logsoftmax_backward(out, out.grad):
    # 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.
 A: $
\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} \\
}$$
