Repository: tensorflow
Issue: Exit code 137 in `tf.raw_ops.ResizeNearestNeighborGrad`
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.16

### 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

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

When the value in the input `size` is too large, program exceeds the memory limit.

### Standalone code to reproduce the issue

```shell
# Signal --4;2024-05-14 00:40:39.895950: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.2024-05-14 00:40:39.896316: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.2024-05-14 00:40:39.900744: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.2024-05-14 00:40:39.956390: 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 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.2024-05-14 00:40:40.842339: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

# ResizeNearestNeighborOpGrad

import tensorflow as tf

align_corners = True
half_pixel_centers = False
grads = tf.constant(1.5e+300, shape=[1,8,16,3], dtype=tf.float64)
size = tf.constant([65534,65534], shape=[2], dtype=tf.int32)
tf.raw_ops.ResizeNearestNeighborGrad(grads=grads, size=size, align_corners=align_corners, half_pixel_centers=half_pixel_centers)
```


### Relevant log output

```shell
(tensorflow-2.16.1-orig) root@b29bda27c601:/mnt# python tests/ResizeNearestNeighborGrad.py 
2024-05-15 08:43:01.439844: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-05-15 08:43:01.440211: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-05-15 08:43:01.444264: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-05-15 08:43:01.498275: 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 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-05-15 08:43:02.744313: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Killed
```
```