Aborted (core dumped) in `tf.raw_ops.NearestNeighbors`
Hi **@LongZE666** ,

- I tried to run your code on colab using TF v2.16.1, nightly and faced the same issue. Please find the [gist](https://colab.sandbox.google.com/gist/Venkat6871/922b66996b0dc68d2a781f9dac3ea427/66765_2-16-nighlty-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.1

### Custom code

Yes

### OS platform and distribution

Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Parameter k being set to a negative number will cause a crash.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

points = tf.constant(0.554979503, shape=[1000,100], dtype=tf.float32)
centers = tf.constant(0.554979503, shape=[1000,100], dtype=tf.float32)
k = tf.constant(-1250999896764, shape=[], dtype=tf.int64)
tf.raw_ops.NearestNeighbors(points=points, centers=centers, k=k)
```


### Relevant log output

```shell
2024-05-01 03:49:00.796702: F tensorflow/core/framework/tensor_shape.cc:201] Non-OK-status: InitDims(dim_sizes) status: INVALID_ARGUMENT: Expected shape dimensions to be non-negative, got -1250999896764
Aborted (core dumped)
```
