Inconsistency in XLA Compiled TensorFlow Model with `tf.nn.softmax, tf.math.reduce_sum` with Extra Transposed Output
Hi @GwiHwan-Go ,

I have replicated the reported behaviour with `jit_compile=True`. With `jit_compile=False` the results are same. Attaching [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/9f8b260d81cb22802a436e44ec6129c0/62519.ipynb) for reference. Needs to check for this behaviour.
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

2.15

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04.3 LTS

### Mobile device

_No response_

### Python version

3.10.0

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

In TensorFlow 2.15.0, a model combining `tf.nn.softmax`, `tf.math.reduce_sum`, and `tf.cast`, with an additional transposed output node, exhibits inconsistent results under XLA compilation compared to eager execution. This inconsistency does not arise when any one of these operations is omitted from the model. This error is seen both on **cpu and gpu.**


### Standalone code to reproduce the issue

```python
import tensorflow as tf
import numpy as np
class Model1(tf.keras.Model):
    @tf.function(jit_compile=True)
    def __call__(self, inp1):
        softmax = tf.nn.softmax(inp1, axis=0)
        trans = tf.transpose(softmax, perm=[0, 2, 1])
        reduce_sum = tf.math.reduce_sum(trans, axis=0)
        cast = tf.cast(reduce_sum, dtype=tf.int64)
        return reduce_sum, cast,

class Model2(tf.keras.Model):
    @tf.function(jit_compile=True)
    def __call__(self, inp1):
        softmax = tf.nn.softmax(inp1, axis=0)
        trans = tf.transpose(softmax, perm=[0, 2, 1])
        reduce_sum = tf.math.reduce_sum(trans, axis=0)
        cast = tf.cast(reduce_sum, dtype=tf.int64)
        return reduce_sum, cast, trans

inputs = [
tf.random.uniform(shape=[5, 5, 5], dtype=tf.float32),
]
model1 = Model1()
model2 = Model2()
device = "cpu"
with tf.device(device):
    tf.config.run_functions_eagerly(True)
    out1 = model1(*inputs)
    out2 = model2(*inputs)
    print(f'=========eager_output(version:{tf.__version__})================')
    try :
        for i in range(min(len(out1),len(out2))):
            np.testing.assert_allclose(out1[i].numpy(), out2[i].numpy(), rtol=0.001, atol=0.001, err_msg=f'at checking {i}th')
        print("XLA_eager does not trigger assertion")
    except AssertionError as e:
        print("XLA_eager triggers assertion")
        print(e)
    tf.config.run_functions_eagerly(False)
    out1 = model1(*inputs)
    out2 = model2(*inputs)
    print(f'=========compiled_output(version:{tf.__version__})================')
    try :
        for i in range(min(len(out1),len(out2))):
            np.testing.assert_allclose(out1[i].numpy(), out2[i].numpy(), rtol=0.001, atol=0.001, err_msg=f'at checking {i}th')
        print("XLA_complie does not trigger assertion")
    except AssertionError as e:
        print("XLA_complie triggers assertion")
        print(e)
```


### Relevant log output

```shell
=========eager_output(version:2.15.0)================
XLA_eager does not trigger assertion
2023-11-30 18:20:25.532764: I external/local_xla/xla/service/service.cc:168] XLA service 0x556187224ad0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2023-11-30 18:20:25.532782: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2023-11-30 18:20:25.535884: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1701368425.554320 2905470 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-11-30 18:20:25.555056: E external/local_xla/xla/stream_executor/stream_executor_internal.h:177] SetPriority unimplemented for this stream.
=========compiled_output(version:2.15.0)================
XLA_complie triggers assertion

Not equal to tolerance rtol=0.001, atol=0.001
at checking 1th
Mismatched elements: 3 / 25 (12%)
Max absolute difference: 1
Max relative difference: 0.
 x: array([[0, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],...
 y: array([[0, 0, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1],...
```
