Different Behavior of tf.raw_ops.BatchMatMulV2+tf.raw_ops.Cos with jit_compile=True
Hello, @zoux1a! 
I was able to replicate the issue reported [here](https://colab.research.google.com/gist/sushreebarsa/3de02546d56ee135a3083b7ae0e83493/62293.ipynb#scrollTo=FZK9eLBV9D6p).
Thank you! Hi,

I see it fails in both` jit_compile = True` and `jit_compile = False`  on GPU, since it invloves random values as an input, results are not guaranteed to match in both the scenarios.
Here is the attached Gist for reference https://gist.github.com/sachinprasadhs/f0d7227796f7dcb74fb5ca735b84f06c This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. 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/62293">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/62293">No</a>

### Issue type

Bug

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

No

### Source

source

### TensorFlow version

2.14.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04.3 LTS (x86_64)

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

GPU 0: NVIDIA GeForce RTX 2070 GPU 1: NVIDIA GeForce RTX 2070 GPU 2: NVIDIA GeForce RTX 2070 GPU 3: NVIDIA GeForce RTX 2070

### Current behavior?

When the **tf.raw_ops.BatchMatMulV2+tf.raw_ops.Cos** operation is invoked within a tf.function with JIT compilation enabled (**jit_compile=True**), it produces different results compared to the same operation called without JIT compilation. This inconsistency is observed when the code is executed on a **GPU** device.
The problem occurs when input Tensors pass through tf.raw_ops.BatchMatMulV2+tf.raw_ops.Cos. With individual Ops there is no issue.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import traceback

class Network(tf.Module):
    def __init__(self):
        super().__init__()

    @tf.function(jit_compile=True)
    def __call__(self, x):
      
      x = tf.raw_ops.BatchMatMulV2(y=x, adj_x=False,adj_y=False,x=tf.random.normal([8, 3], dtype=tf.float32))        
      x = tf.raw_ops.Cos(x=x, )        
      return x

m = Network()
inp = {
    "x": tf.random.normal([8, 3, 8, 7, 3, 4], dtype=tf.float32),
}

with tf.device('/GPU:0'):
    tf.config.run_functions_eagerly(True)
    no_op_res = m(**inp)
    tf.config.run_functions_eagerly(False)
    with tf.device('/GPU:0'):
        op_res = m(**inp)

    tf.debugging.assert_near(tf.cast(no_op_res, tf.float64), tf.cast(op_res, tf.float64), atol=0.001, rtol=0.001)
```


### Relevant log output

_No response_