Repository: tensorflow
Issue: `log_softmax` could be `2**102` to `2**970` times more accurate
Labels: stat:awaiting tensorflower,type:feature,comp:apis

### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

unknown 2.12.0

### Custom code

No

### OS platform and distribution

Linux 5.15.0-67-generic #74~20.04.1-Ubuntu SMP Wed Feb 22 14:52:34 UTC 2023 x86_64 GNU/Linux

### Mobile device

_No response_

### Python version

3.11.4

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

As I understand it, TensorFlow implements `log_softmax(x)` as `x - max(x) - log(sum(exp(x - max(x))))`. 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 by using `log1p` and zeroing out the maximum cell (guaranteed to be 1) before summing.

This came up when a transformer I was training with cross-entropy loss on a classification task had loss dominated by float32 error.

Might help with https://github.com/tensorflow/tensorflow/issues/48816, https://github.com/tensorflow/tensorflow/issues/12002

I originally posted this as a [StackOverflow question](https://stackoverflow.com/q/77477327/377022).
Companion PyTorch issue: https://github.com/pytorch/pytorch/issues/113708
Companion SciPy issue: https://github.com/scipy/scipy/issues/19521

### Standalone code to reproduce the issue

```python
import tensorflow as tf
eps = tf.constant(tf.keras.backend.epsilon(), dtype=tf.float32)
print(-tf.nn.log_softmax(tf.constant([1 - tf.math.log(2 * eps).numpy(), 0], dtype=tf.float32))) # tf.Tensor([1.1920928e-07 1.6424950e+01], shape=(2,), dtype=float32)
print(-tf.nn.log_softmax(tf.constant([1 - tf.math.log(eps).numpy(), 0], dtype=tf.float32))) # tf.Tensor([-0.       17.118095], shape=(2,), dtype=float32)
```

```python
import tensorflow as tf
import numpy as np

def log_softmax_alt(x):
    maxi = tf.argmax(x)
    xoffset = x - x[maxi]
    xoffsetexp = tf.exp(xoffset)
    # xoffsetexp[maxi] is currently about 1
    xoffsetexp = tf.tensor_scatter_nd_update(xoffsetexp, [[maxi]], [0])
    xoffsetexp_sum_m1 = tf.reduce_sum(xoffsetexp)
    return xoffset - tf.math.log1p(xoffsetexp_sum_m1)


for ty in (tf.float32, tf.float64):
    smallest_log_softmax, smallest_log_softmax_alt, smallest_log_softmax_val, smallest_log_softmax_alt_val = 0, 0, 0, 0
    # work around https://github.com/tensorflow/tensorflow/issues/62399
    for i in range(int(1 - np.log2(tf.experimental.numpy.finfo(ty).smallest_subnormal))):
        values = tf.constant([1 + i, 0], dtype=ty)
        log_softmax_values = tf.nn.log_softmax(values)
        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]
        if log_softmax_values_alt[0] != 0: smallest_log_softmax_alt, smallest_log_softmax_alt_val = i, log_softmax_values_alt[0]
        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)}")
```
gives
```
For <dtype: 'float32'>, diff in supported input accuracy is 2**-(86 - 15) = 2**-71; diff in output accuracy is np.log2(1.1920928244535389e-07) - np.log2(1.6458114537543937e-38) = 102.51446533203125
For <dtype: 'float64'>, diff in supported input accuracy is 2**-(707 - 35) = 2**-672; diff in output accuracy is np.log2(2.2204460492503128e-16) - np.log2(3.3075530036384083e-308) = 969.428088949386
```


### Relevant log output

_No response_