JIT yields inconsistent results using tf.math.top_k
@nicolaspi Could you please have a look at the colab gist of TF [v2.13](https://colab.research.google.com/gist/sushreebarsa/eec8d03e8c0fd64ca06dd633ffa31057/untitled840.ipynb#scrollTo=GE1XkYdiU6ie) and tf-[nightly](https://colab.research.google.com/gist/sushreebarsa/aec2f99123ca0e0b9b4268a14995cc12/untitled840.ipynb#scrollTo=OvarIK-iTOX2). I couldn't reproduce the error reported here. 
The output is as follows;
```
should_be_all_zero shape (1024, 3):
[[0.]
 [0.]
 [0.]]
should_be_all_zero shape (1023, 3):
[[0.]
 [0.]
 [0.]]
should_be_all_zero shape (1024, 2):
[[0.]
 [0.]]
```
Thank you!  Updating the GPU drivers fixed the issue.
Thanks Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61692">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61692">No</a>
 @nicolaspi Glad your issue has been fixed. 
Thank you!
### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

2.15.0-dev20230824

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

cuda_11.8.r11.8/compiler.31833905_0 / cuDNN version 8700

### GPU model and memory

NVIDIA GeForce RTX 2080 Ti

### Current behavior?

JIT yields inconsistent results using `tf.math.top_k` when `index_type=tf.int32` (no issue with `index_type=tf.int64`).

### Standalone code to reproduce the issue

```shell
import tensorflow as tf


@tf.function(jit_compile=True)
def tf_func(shape):
    x = tf.random.stateless_normal(shape, seed=(1, 2))
    x = tf.transpose(x, perm=[1, 0])
    topk_max, indices = tf.math.top_k(x, 1, sorted=False, index_type=tf.int32)
    reduce_max = tf.reduce_max(x, axis=1, keepdims=True)
    return topk_max - reduce_max


def check(shape):
    should_be_all_zero = tf_func(shape)
    print(f"should_be_all_zero shape {shape}:\n{should_be_all_zero}")


check((1024, 3))
check((1023, 3))
check((1024, 2))
```


### Relevant log output

```shell
should_be_all_zero shape (1024, 3):
[[-0.7787831 ]
 [ 0.47324872]
 [ 0.30553436]]
should_be_all_zero shape (1023, 3):
[[0.]
 [0.]
 [0.]]
should_be_all_zero shape (1024, 2):
[[0.]
 [0.]]
```
