Softmax That Does Not Overflow

Every transformer runs this in its inner loop, and the naive version silently produces nan on real logits.

Implement softmax(x, axis=-1) for an array of any shape:

softmax(x)[i] = exp(x[i]) / sum(exp(x))

along the given axis.

The trap: exp(1000) is inf, and inf / inf is nan. Logits of that magnitude are ordinary in an untrained or badly scaled model, and a loss that turns to nan three hours into training is a genuinely expensive bug.

Your implementation must return correct probabilities for inputs where a naive np.exp(x) overflows.