System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below):2.1.0
• Python version:3.7.6
• Bazel version (if compiling from source):N/A
• GCC/Compiler version (if compiling from source):N/A
• CUDA/cuDNN version:N/A
• GPU model and memory:N/A

Describe the current behavior
tf.math.segment_min abortion when segment_ids contains large value
Describe the expected behavior
expect an exception message if the input is not expected, instead of crash.

Standalone code to reproduce the issue

import tensorflow as tf
import numpy as np
tf.math.segment_min(data=np.ones((1,2,1,1,1), dtype=np.uint16), segment_ids=[5053997376933981534])

output:

2021-01-26 15:49:10.870684: F tensorflow/core/framework/tensor_shape.cc:405] Check failed: 0 <= new_num_elements (0 vs. -8338749319841588546)
Aborted (core dumped)