Aborted (core dumped) in `tf.raw_ops.ResourceApplyRMSProp/tf.raw_ops.ResourceSparseApplyRMSProp`
@x0w3n I was able to replicate the issue reported here using TF [v2.16](https://colab.research.google.com/gist/sushreebarsa/e83f09c20b67c474866a7a6d165ac835/69281.ipynb#scrollTo=7OK58fNuo1mD) and tf-[nightly](https://colab.research.google.com/gist/sushreebarsa/24a762d4f3f5feda5c4a69563f4e0347/69281.ipynb).
Thank you!
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.16

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04.3 LTS (x86_64)

### Mobile device

_No response_

### Python version

3.9.13

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

On specific input, tf.raw_ops.ResourceApplyRMSProp/tf.raw_ops.ResourceSparseApplyRMSProp triggers "Aborted (core dumped)". We ran the code on colab's latest TensorFlow Nightly, which also triggers the crash. The cause of the crash may be the grad parameter.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

var = tf.Variable([1.0, 2.0, 3.0])
accum = tf.Variable([0.1, 0.2, 0.3], dtype=tf.complex64) 
accum2 = tf.Variable([0, 1.2, -3.3], dtype=tf.complex64) 
lr = 0.01
grad = tf.constant([0.1, 0.2, 0.3])
momentum = 0.9

# crash
tf.raw_ops.ResourceApplyRMSProp(
    var=var.handle,
    ms=accum.handle,
    mom=accum2.handle,
    lr=lr,
    rho=lr,
    momentum=momentum,
    epsilon=0.9,
    grad=grad,
    use_locking=False,
)

# tf.raw_ops.ResourceSparseApplyRMSProp(
#     var=var.handle,
#     ms=accum.handle,
#     mom=accum2.handle,
#     lr=lr,
#     rho=lr,
#     momentum=momentum,
#     indices=tf.constant([1,2,2]),
#     epsilon=0.9,
#     grad=grad,
#     use_locking=False,
# )
```


### Relevant log output

```shell
2024-06-06 01:07:00.757682: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-06 01:07:01.835541: F tensorflow/core/framework/tensor.cc:844] Check failed: dtype() == expected_dtype (8 vs. 1) float expected, got complex64
Aborted (core dumped)
```
