tf.sparse.segment_sqrt_n/sum crashes(abort) when segment_ids is large uint

[DNXie]

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.sparse.segment_sqrt_n and tf.sparse.segment_sum crashes(abort) when segment_ids is large uint

Describe the expected behavior
expect no crash if the input is unexpected instead of crash

Standalone code to reproduce the issue

tf.sparse.segment_sqrt_n(data=1.0, indices=[1], segment_ids=np.array([1205444461],dtype=np.uint32))
tf.sparse.segment_sum(data=1, indices=[1], segment_ids=np.array([1205444461], dtype=np.uint32))

Output:

2021-02-03 21:19:17.206426: F tensorflow/core/framework/tensor_shape.cc:435] Check failed: d < dims() (0 vs. 0)
Aborted (core dumped)

[Saduf2019]

@ymodak

I ran the code shared on tf 2.4 and nightly, colab crashes. please find the gist here

[DNXie]

Similar issue in tf.sparse.segment_mean too

tf.sparse.segment_mean(data=1.0, indices=[1], segment_ids=np.array([1205444461],dtype=np.uint32))

Output:

Aborted (core dumped)

[rainwoodman]

 tf.sparse.segment_sum(data=1.0, indices=[0], segment_ids=np.array([3],dtype=np.uint32))

Also crashes.

But

tf.sparse.segment_sum(data=[1.0], indices=[0], segment_ids=np.array([1205444461],dtype=np.uint32))

Works.

So the problem is not caused large id, but it is likely a memory allocation error during the error handling in the eager mode when the shape of data is incorrect.

The function does not consume a scalar for the data argument.

@tf.function
def func():
  return tf.sparse.segment_sum(data=1.0, indices=[0], segment_ids=np.array([1205444461],dtype=np.uint32))
func()

reports the shape error correctly.

[sushreebarsa]

Was able to reproduce the issue in TF v2.5,Please find the gist here..Thanks !

[rainwoodman]

Thanks! I will try to push a fix this week.

[ymodak]

This is fixed with tf-nightly and will be part of TF 2.6 release.
Please see gist for reference. Thanks!