Merge pull request #51359 from yongtang:46913-range-overflow

PiperOrigin-RevId: 391529518
Change-Id: Ie3db4ae6d3c0f3dc88404e1dbdc22f7d03cbeb3b

tensorflow/core/kernels/sequence_ops.cc

@@ -71,12 +71,13 @@ class RangeOp : public OpKernel {
           errors::InvalidArgument(
               "Requires start >= limit when delta < 0: ", start, "/", limit));
     }
-    int64 size = (std::is_integral<T>::value
-                      ? ((std::abs(limit - start) + std::abs(delta) - 1) /
-                         std::abs(delta))
-                      : std::ceil(std::abs((limit - start) / delta)));
-    TensorShape shape;
-    OP_REQUIRES_OK(context, shape.AddDimWithStatus(size));
+    int64 size = 0;
+    if (std::is_integral<T>::value) {
+      size = static_cast<int64>(
+          (std::abs(limit - start) + std::abs(delta) - 1) / std::abs(delta));
+    } else {
+      size = static_cast<int64>(std::ceil(std::abs((limit - start) / delta)));
+    }
     Tensor* out = nullptr;
     OP_REQUIRES_OK(context, context->allocate_output(0, shape, &out));
     auto flat = out->flat<T>();

tensorflow/python/kernel_tests/init_ops_test.py

@@ -23,6 +23,7 @@
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import test_util
@@ -541,13 +542,24 @@ def testMixedDType(self):
     tf_ans = math_ops.range(
         constant_op.constant(4, dtype=dtypes.int32), dtype=dtypes.int64)
     self.assertAllEqual(self.evaluate(tf_ans), np.array([0, 1, 2, 3]))
+<<<<<<< HEAD
 
  def testLargeStarts(self):
     # Test case for GitHub issue 46899.
     with self.session():
       with self.assertRaises(errors_impl.InternalError):
         v = math_ops.range(start=-1e+38, limit=1)
         self.evaluate(v)
+=======
+
+  def testLargeLimits(self):
+    # Test case for GitHub issue 46913.
+    with self.session():
+      with self.assertRaises(errors_impl.ResourceExhaustedError):
+        v = math_ops.range(0, 9223372036854775807)
+        self.evaluate(v)
+
+>>>>>>> 6d94002a097 (Merge pull request #51359 from yongtang:46913-range-overflow)
 
 # TODO(vrv): move to sequence_ops_test?
 class LinSpaceTest(test.TestCase):