Merge pull request #51733 from yongtang:46888-tf.math.segment_

PiperOrigin-RevId: 406020083
Change-Id: I179a9a8fe548ed324fc97363e81a46be28aa19b8

tensorflow/core/kernels/segment_reduction_ops_impl.h

@@ -110,7 +110,9 @@ class SegmentReductionOp : public OpKernel {
                 errors::InvalidArgument("Shape must be at least rank 1"));
 
     TensorShape output_shape = input.shape();
-    output_shape.set_dim(0, output_rows);
+    // Since we're changing the first dimension of the shape, we need to make
+    // sure the new shape won't overflow.
+    OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, output_rows));
 
     // Note that we do not initialize the output buffer with a default value, so
     // we need to explicitly set missing indices to the default value.
@@ -290,7 +292,10 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
                         done);
 
       TensorShape output_shape = input.shape();
-      output_shape.set_dim(0, output_rows);
+      // Since we're changing the first dimension of the shape, we need to make
+      // sure the new shape won't overflow.
+      OP_REQUIRES_OK_ASYNC(context,
+                           output_shape.SetDimWithStatus(0, output_rows), done);
 
       Tensor* output = nullptr;
       OP_REQUIRES_OK_ASYNC(

tensorflow/python/kernel_tests/segment_reduction_ops_test.py

@@ -270,6 +270,20 @@ def testDataInvalid(self):
               data=np.uint16(10), segment_ids=np.array([]).astype("int64"))
           self.evaluate(s)
 
+  def testInvalidIds(self):
+    # Test case for GitHub issue 46888.
+    for op in [
+        math_ops.segment_max,
+        math_ops.segment_min,
+        math_ops.segment_mean,
+        math_ops.segment_sum,
+        math_ops.segment_prod,
+    ]:
+      with self.cached_session(use_gpu=False):
+        with self.assertRaises((ValueError, errors_impl.InternalError)):
+          s = op(data=np.ones((1, 10, 1)), segment_ids=[1676240524292489355])
+          self.evaluate(s)
+
 
 class UnsortedSegmentTest(SegmentReductionHelper):