A Fix for Softmax and Cross-Entropy

The authors show that the standard log-softmax and cross-entropy implementations shipped by JAX, PyTorch and TensorFlow contain two numerical bugs they call Softmax Collapse (the target logit's gradient rounds to zero when it dominates, so gradients no longer sum to zero) and Softmax Cancellation (non-zero gradients when student and teacher logits are equal). Drop-in corrected implementations eliminate the known 'Slingshot' phenomenon, speed up grokking with weight decay, prevent accuracy collapse without weight decay, and stop distillation students drifting from the teacher.