Problem with chain rule in softmax layer when differentiated separately
I'm just gonna answer my own question here so that in the future I can find it easier, and hopefully it comes handy for others asking about the derivative of Softmax and gets confused by the fact that it's a matrix. So the shown matrix form for $\frac{\partial{Softmax(x)}}{\partial{x}} $ in the opening question is incorrect. The correct formula should be:
$$ \frac{\partial{Softmax(Z^{O})}}{\partial{Z^{O}}} := \begin{pmatrix}\frac{\partial{Softmax(Z^{O}_{0})}}{\partial{Z^{O}_{0}}}&amp;\frac{\partial{Softmax(Z^{O}_{0})}}{\partial{Z^{O}_{1}}}&amp;\frac{\partial{Softmax(Z^{O}_{0})}}{\partial{Z^{O}_{2}}}&amp;...&amp;\frac{\partial{Softmax(Z^{O}_{0})}}{\partial{Z^{O}_{N}}}\\\frac{\partial{Softmax(Z^{O}_{1})}}{\partial{Z^{O}_{0}}}&amp;\frac{\partial{Softmax(Z^{O}_{1})}}{\partial{Z^{O}_{1}}}&amp;\frac{\partial{Softmax(Z^{O}_{1})}}{\partial{Z^{O}_{2}}}&amp;...&amp;\frac{\partial{Softmax(Z^{O}_{1})}}{\partial{Z^{O}_{N}}}\\\frac{\partial{Softmax(Z^{O}_{2})}}{\partial{Z^{O}_{0}}}&amp;\frac{\partial{Softmax(Z^{O}_{2})}}{\partial{Z^{O}_{1}}}&amp;\frac{\partial{Softmax(Z^{O}_{2})}}{\partial{Z^{O}_{2}}}&amp;...&amp;\frac{\partial{Softmax(Z^{O}_{2})}}{\partial{Z^{O}_{N}}}\\...&amp;...&amp;...&amp;&amp;...\\\frac{\partial{Softmax(Z^{O}_{N})}}{\partial{Z^{O}_{0}}}&amp;\frac{\partial{Softmax(Z^{O}_{N})}}{\partial{Z^{O}_{1}}}&amp;\frac{\partial{Softmax(Z^{O}_{N})}}{\partial{Z^{O}_{2}}}&amp;...&amp;\frac{\partial{Softmax(Z^{O}_{N})}}{\partial{Z^{O}_{N}}}\end{pmatrix}$$
Which is:
$$ \begin{pmatrix}S(Z^{O}_{0})(1-S(Z^{O}_{0}))&amp;-S(Z^{O}_{0})S(Z^{O}_{1})&amp;-S(Z^{O}_{0})S(Z^{O}_{2})&amp;...&amp;-S(Z^{O}_{0})S(Z^{O}_{N})\\-S(Z^{O}_{1})S(Z^{O}_{0})&amp;S(Z^{O}_{1})(1-S(Z^{O}_{1}))&amp;-S(Z^{1}_{1})S(Z^{O}_{2})&amp;...&amp;-S(Z^{O}_{1})S(Z^{O}_{N})\\-S(Z^{O}_{2})S(Z^{O}_{0})&amp;-S(Z^{O}_{2})S(Z^{O}_{1})&amp;S(Z^{2}_{2})(1-S(Z^{2}_{2}))&amp;...&amp;-S(Z^{O}_{2})S(Z^{O}_{N})\\...&amp;...&amp;...&amp;&amp;...\\-S(Z^{O}_{N})S(Z^{O}_{0})&amp;-S(Z^{O}_{N})S(Z^{O}_{1})&amp;-S(Z^{O}_{N})S(Z^{O}_{2})&amp;...&amp;S(Z^{O}_{N})(1-S(Z^{O}_{N}))\end{pmatrix} $$
For which I incorrectly used the df2 function with z1 argument, as it should be:
$$ \begin{pmatrix}A^O_0(1-A^O_0)&amp;-A^O_0A^O_1\\-A^O_0A^O_1&amp;A^O_1(1-A^O_1)\end{pmatrix} $$
Using this, I have:
$$ \frac{\partial{A^{O}}}{\partial{Z^{O}}} = Softmax^{'}(A^{O}) = \begin{pmatrix}.25&amp;-.25\\-.25&amp;.25\end{pmatrix} $$
Finally doing the multiplication:
$$ \frac{\partial{E(A^{O}, T)}}{\partial{A^{O}}} * \frac{\partial{A^{O}}}{\partial{Z^{O}}} = \begin{pmatrix}0&amp;-2\end{pmatrix}*\begin{pmatrix}0.25&amp;-0.25\\-0.25&amp;0.25\end{pmatrix} = \begin{pmatrix}0.5&amp;-0.5\end{pmatrix} $$
Which matches the original:
$$ \frac{\partial{E(A^{O}, T)}}{\partial{Z^{O}}} = A^{O} - T = \begin{pmatrix}0.5&amp;0.5\end{pmatrix} - \begin{pmatrix}0&amp;1\end{pmatrix} = \begin{pmatrix}0.5&amp;-0.5\end{pmatrix} $$
And this is where I messed up.