Repository: pytorch
Issue: `log_softmax` could be `2**124` to `2**1021` times more accurate on small outputs
Labels: module: nn,triaged

### 🚀 The feature, motivation and pitch

As I understand it, PyTorch implements `log_softmax(x)` as `x - x.max() - (x - x.max()).exp().sum().log()`. However, when the largest value is much larger than the rest of the values (about 16 larger for float32, about 36 larger for float64), `log_softmax` returns `0` at the maximum value, when it could give a much more precise answer.

Numerically, we have
```python
>>> eps = torch.tensor(torch.finfo(torch.float32).eps, dtype=torch.float32)
>>> -torch.tensor([1-(2*eps).log(), 0]).log_softmax(dim=0)
tensor([1.1921e-07, 1.6249e+01])
>>> -torch.tensor([1-eps.log(), 0]).log_softmax(dim=0)
tensor([-0.0000, 16.9424])
```

(For float16, this is 7.2e-4, for float64 2.2e-16.)

The value of `(x - x.max()).exp().sum()` is going to be at least 1, because the largest value of `x - x.max()` is 0.  1.19e-7 is the unit of least precision for float32 around 1, but float32 can represent values as small as about 1.18e-38 normally (`torch.finfo(torch.float32).smallest_normal`, `2**-126 = 2**-(2**7 - 2)`) or about 1.18e-45 (`2**-126 * 2**-23`) subnormally.  So there's a lot of room for more precision using `log1p(x) = log(1 + x)` instead.  Consider:
```python
eps = torch.tensor(torch.finfo(torch.float32).smallest_normal, dtype=torch.float32)
x = torch.tensor([1-eps.log(), 0])
maxi = x.argmax()
xoffset = x - x[maxi]
xoffsetexp = xoffset.exp()
# xoffsetexp[maxi] is currently about 1
xoffsetexp[maxi] = 0
xoffsetexp_sum_m1 = xoffsetexp.sum()
print(-(xoffset - xoffsetexp_sum_m1.log1p())) # tensor([4.3244e-39, 8.8337e+01])
print((-(xoffset - xoffsetexp_sum_m1.log1p())[0]).log2()) # tensor(-127.4427)
```

For float64 and using `smallest_subnormal` instead of `smallest_normal`, the difference is even more pronounced:
```python
eps = torch.tensor(numpy.finfo(numpy.float64).smallest_subnormal, dtype=torch.float64)
x = torch.tensor([1-(8*eps).log(), 0])
maxi = x.argmax()
xoffset = x - x[maxi]
xoffsetexp = xoffset.exp()
# xoffsetexp[maxi] is current about 1
xoffsetexp[maxi] = 0
xoffsetexp_sum_m1 = xoffsetexp.sum()
print(-(xoffset - xoffsetexp_sum_m1.log1p())) # tensor([1.9763e-323,  7.4336e+02], dtype=torch.float64)
print((-(xoffset - xoffsetexp_sum_m1.log1p())[0]).log2()) # tensor(-1072., dtype=torch.float64)
```

This came up when a transformer I was training with cross-entropy loss on a classification task had loss dominated by `torch.finfo(torch.float32).eps`.

### Alternatives

_No response_

### Additional context

Numbers in the title were generated with
```python
import torch as t
import numpy as np

def log_softmax_alt(x):
    maxi = t.argmax(x)
    xoffset = x - x[maxi]
    xoffsetexp = t.exp(xoffset)
    # xoffsetexp[maxi] is currently about 1
    xoffsetexp[maxi] = 0
    xoffsetexp_sum_m1 = t.sum(xoffsetexp)
    return xoffset - t.log1p(xoffsetexp_sum_m1)


for ty, nty in ((t.float32, np.float32), (t.float64, np.float64)):
    smallest_log_softmax, smallest_log_softmax_alt, smallest_log_softmax_val, smallest_log_softmax_alt_val = 0, 0, 0, 0
    for i in range(int(1-t.log2(t.tensor(np.finfo(nty).smallest_subnormal, dtype=ty)))):
        values = t.tensor([1 + i, 0], dtype=ty)
        log_softmax_values = values.log_softmax(dim=0)
        log_softmax_values_alt = log_softmax_alt(values)
        if log_softmax_values[0] != 0: smallest_log_softmax, smallest_log_softmax_val = i, log_softmax_values[0].item()
        if log_softmax_values_alt[0] != 0: smallest_log_softmax_alt, smallest_log_softmax_alt_val = i, log_softmax_values_alt[0].item()
        if log_softmax_values[0] == 0 and log_softmax_values_alt[0] == 0: break
    print(f"For {ty}, diff in supported input accuracy is 2**-({smallest_log_softmax_alt} - {smallest_log_softmax}) = 2**-{smallest_log_softmax_alt - smallest_log_softmax}; diff in output accuracy is np.log2({-smallest_log_softmax_val}) - np.log2({-smallest_log_softmax_alt_val}) = {np.log2(-smallest_log_softmax_val) - np.log2(-smallest_log_softmax_alt_val)}")
```
giving
For torch.float32, diff in supported input accuracy is 2**-(101 - 15) = 2**-86; diff in output accuracy is np.log2(1.1920928244535389e-07) - np.log2(5.605193857299268e-45) = 123.99999991400867
For torch.float64, diff in supported input accuracy is 2**-(743 - 35) = 2**-708; diff in output accuracy is np.log2(2.2204460492503128e-16) - np.log2(1e-323) = 1021.0

I originally posted this as a [StackOverflow question](https://stackoverflow.com/q/77477327/377022).
Companion SciPy issue: https://github.com/scipy/scipy/issues/19521
Companion TensorFlow issue: https://github.com/tensorflow/tensorflow/issues/62400

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki