Softmax overflow issue with a large tensor
related to https://github.com/tensorflow/tensorflow/issues/60330 @sachinprasadhs,
I was able to reproduce the issue on tensorflow v2.12, v2.13 and tf-nightly. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/5ba9acb20cc199769dc8cc79ab685fe2/61639.ipynb). any updates? thanks
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

tf 2.14.0 dev20230706

### Custom code

Yes

### OS platform and distribution

Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

A100 80G

### Current behavior?

Non-OK-status: GpuLaunchKernel( GenerateNormalizedProb<T, acc_type, kUnroll>, numBlocks, numThreadsPerBlock, 0, cu_stream, reinterpret_cast<const T*>(logits_in_.flat<T>().data()), reinterpret_cast<const acc_type*>( sum_probs.flat<acc_type>().data()), reinterpret_cast<const T*>(max_logits.flat<T>().data()), const_cast<T*>(softmax_out->flat<T>().data()), rows, cols, log_) status: INTERNAL: invalid configuration argument

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import numpy as np

print(tf.__version__)

x = tf.constant(1., shape=(1, 1))
y = tf.tile(x, (2**28, 9))  # Number of elements cannot fit in int32 value
print(y.numpy())
z = tf.math.softmax(y)
print(z.numpy())
```


### Relevant log output

```shell
2023-08-18 13:34:53.496914: I tensorflow/core/common_runtime/placer.cc:125] logits: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
Softmax: (Softmax): /job:localhost/replica:0/task:0/device:GPU:0
2023-08-18 13:34:53.496979: I tensorflow/core/common_runtime/placer.cc:125] Softmax: (Softmax): /job:localhost/replica:0/task:0/device:GPU:0
softmax_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-08-18 13:34:53.496988: I tensorflow/core/common_runtime/placer.cc:125] softmax_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-08-18 13:34:53.497454: I tensorflow/core/common_runtime/eager/execute.cc:1747] Executing op Softmax in device /job:localhost/replica:0/task:0/device:GPU:0
2023-08-18 13:34:53.507131: F tensorflow/core/kernels/softmax_op_gpu.cu.cc:255] Non-OK-status: GpuLaunchKernel( GenerateNormalizedProb<T, acc_type, kUnroll>, numBlocks, numThreadsPerBlock, 0, cu_stream, reinterpret_cast<const T*>(logits_in_.flat<T>().data()), reinterpret_cast<const acc_type*>( sum_probs.flat<acc_type>().data()), reinterpret_cast<const T*>(max_logits.flat<T>().data()), const_cast<T*>(softmax_out->flat<T>().data()), rows, cols, log_) status: INTERNAL: invalid configuration argument
Aborted (core dumped)
```
