`Check failed` in `tf.strings.unsorted_segment_join` when the rank of `inputs` is larger than 1 and the value of `num_segments` is too large.
Hi **@xl1043237213** ,

I tried to run your code on colab using TF v2.15, 2.16.1, nightly and faced the same issue. Please find the [gist](https://colab.sandbox.google.com/gist/Venkat6871/ffd39afa22b162d49b6a4475728a3145/65632_2-15-2-16-nightlty-v.ipynb) here for reference.

Thank you!
### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

No

### OS platform and distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

`Check failed` in `tf.strings.unsorted_segment_join` when the rank of `inputs` is larger than 1 and the value of `num_segments` is too large, which causes the program to crash.

### Standalone code to reproduce the issue

```shell
https://colab.research.google.com/drive/1y6fNdHnw5hjfOdxTa1ojr0HCwh8nN94c?usp=sharing
```


### Relevant log output

```shell
Check failed: 0 <= new_num_elements (0 vs. -1)
Aborted (core dumped)
```
