Softmax layer diverges from jax.nn.softmax

[martin-gorner]

repro colab (10 lines of code): https://colab.research.google.com/drive/1F9-I8Ax-OBKT0-6GmOq5macFBpu-cFid?usp=sharing

The keras and jax version of softmax have a different behavior on data that is completely masked out.