Shift by the max in log_softmax

[ayaangazali]

Proposed changes

nn.log_softmax loses the normalizer once the logits get large. Two equal logits have to give log(1/2) no matter their magnitude, but:

>>> [nn.log_softmax(mx.array([[v, v]]))[0, 0].item() for v in (1e0, 1e4, 1e8, 1e20)]
[-0.6931471824645996, -0.693359375, 0.0, 0.0]

pytorch returns -0.6931471824645996 for all four. It is already visibly off at 1e4 and reports 0.0 from 1e8 on, which claims each class has probability 1.

The cause is the composition:

return x - mx.logsumexp(x, axis=axis, keepdims=True)

mx.logsumexp itself is fine. logsumexp([1e8, 1e8]) is 1e8 + log(2), and since the ulp of 1e8 in float32 is 8, 1e8 is the correctly rounded answer. The problem is that the normalizer is added to a large magnitude before it is subtracted off again, so it rounds away and never reaches the result.

Shifting by the max first keeps the normalizer at a magnitude where it survives, which is what pytorch does:

x = x - mx.stop_gradient(mx.max(x, axis=axis, keepdims=True))
return x - mx.logsumexp(x, axis=axis, keepdims=True)

After the change every scale from 1e0 to 1e36 gives -0.693147, matching pytorch. On random inputs the two now agree exactly (max difference 0) for scales of 1e3 and above, and both sit the same distance from a float64 reference, which is the float32 limit of the inputs rather than of the algorithm. Gradients are unchanged, max difference against pytorch 0.0, and stop_gradient on the shift keeps the graph the same shift invariant function it was.

Special values line up with pytorch better than before as a side effect: log_softmax([inf, 1]) was [nan, -inf] and is now [nan, nan].

mx.softmax was never affected. It is a C++ op that already subtracts the max, and it returns [0.5, 0.5] at every scale.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

CPU only build (MLX_BUILD_METAL=OFF) at a076a63. test_nn.py, test_losses.py, test_optimizers.py and test_ops.py all pass, and the new assertion fails on main from 1e8 up. Python only, so no rebuild is involved.

---

freshman contributor, i lean on Claude Code while digging. found this feeding extreme values through the activations and comparing against pytorch. one thing i want to flag rather than hide: my first instinct was that mx.logsumexp was wrong, and it took working out the ulp at 1e8 to see that it returns the best float32 answer available and the fault is in how it gets combined here.

[zcbenz]

Thanks!

Cross-checked JAX and it is doing the same:
https://github.com/jax-ml/jax/blob/da2d20de08edc1041b1983d0800d12c733a51352/jax/_src/nn/functions.py#L569