Difference between mathematical and Tensorflow implementation of Softmax Crossentropy with logit
As I understand it, the softmax function for $z_i$ is given by $a_i$. Then just taking the loss you've defined you get back exactly the formula that is implemented. The way it is written down however is, as you mentioned, to avoid underflow/overflow. 

For instance, suppose you want to compute the following:

$A=\log(\sum_{i=1}^{4}\exp(z_i))$, with $z_i=(-1000.5,-2000.5,-3000.5,-4000.5)$

Clearly, if you just type in the formula directly, you will get an underflow error. Instead if you isolate the main contribution in the exponential by taking the $\max(z_i)$, the same formula can be written as:

$A=\max_i(z_i)+\log(\sum_{i=1}^{4}\exp(z_i-\max_i(z_i)))$

The difference now is that the expression is "numerically stable" and we see that $A\approx -1000.5$.

Thus, let's make the softmax numerically stable:
\begin{align}
\log(a_i)&amp;=z_i-\log(\sum_j e^{z_j})\\
&amp;=z_i-\max_j(z_j)-\log(\sum_je^{z_j-\max_j(z_j)})
\end{align}
which is the expression that is implemented for the loss (just multiply by $y_i$ and sum over $i$).