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_max/min/mean/sun/prod crashes(aborts) when segment_ids is large

Describe the expected behavior
expect an exception message if the input is unexpected instead of crash

Standalone code to reproduce the issue

tf.math.segment_max(data=np.ones((1,10,1)), segment_ids=[1676240524292489355])
tf.math.segment_min(data=np.ones((1,10,1)), segment_ids=[1676240524292489355])
tf.math.segment_mean(data=np.ones((1,10,1)), segment_ids=[1676240524292489355])
tf.math.segment_sum(data=np.ones((1,10,1)), segment_ids=[1676240524292489355])
tf.math.segment_prod(data=np.ones((1,10,1)), segment_ids=[1676240524292489355])

Output:

2021-02-03 16:44:25.849065: F tensorflow/core/framework/tensor_shape.cc:405] Check failed: 0 <= new_num_elements (0 vs. -1684338830784658056)
Aborted (core dumped)

Related issue: #46696