Aborted (core dumped) in `tf.raw_ops.ResourceApplyKerasMomentum/tf.raw_ops.ResourceSparseApplyKerasMomentum`
Hi @x0w3n ,
- I was able to reproduce the issue on Colab using TF v2.15 and TF-nightly ,Please find the [gist](https://colab.research.google.com/gist/Venkat6871/007a03f4964b1cbf66b50026411a9f96/69279_2-15-nightly-v.ipynb) here for reference.

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.ResourceApplyKerasMomentum/tf.raw_ops.ResourceSparseApplyKerasMomentum triggers "Aborted (core dumped)". We ran the code on colab's latest TensorFlow Nightly, which also triggers the crash. We analyzed it, and the cause of this crash is probably 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.ResourceApplyKerasMomentum(
    var=var.handle,
    accum=accum.handle,
    grad=grad,
    lr=lr,
    momentum=0.1,
    use_locking=False,
    use_nesterov=False,
    name=None
)

# crash
# tf.raw_ops.ResourceSparseApplyKerasMomentum(
#     var=var.handle,
#     accum=accum.handle,
#     grad=grad,
#     lr=lr,
#     indices=tf.constant([1,2,2]),
#     momentum=0.1,
#     use_locking=False,
#     use_nesterov=False,
#     name=None
# )
```


### Relevant log output

```shell
2024-06-06 00:59:43.994008: 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 00:59:44.992853: F tensorflow/core/framework/tensor.cc:844] Check failed: dtype() == expected_dtype (8 vs. 1) float expected, got complex64
Aborted (core dumped)
```
