Skip to content

Commit 6d94002

Browse files
Merge pull request #51359 from yongtang:46913-range-overflow
PiperOrigin-RevId: 391529518 Change-Id: Ie3db4ae6d3c0f3dc88404e1dbdc22f7d03cbeb3b
2 parents c2cf882 + 89b5fa3 commit 6d94002

File tree

2 files changed

+15
-4
lines changed

2 files changed

+15
-4
lines changed

Diff for: tensorflow/core/kernels/sequence_ops.cc

+7-4
Original file line numberDiff line numberDiff line change
@@ -71,10 +71,13 @@ class RangeOp : public OpKernel {
7171
errors::InvalidArgument(
7272
"Requires start >= limit when delta < 0: ", start, "/", limit));
7373
}
74-
int64_t size = (std::is_integral<T>::value
75-
? ((std::abs(limit - start) + std::abs(delta) - 1) /
76-
std::abs(delta))
77-
: std::ceil(std::abs((limit - start) / delta)));
74+
int64_t size = 0;
75+
if (std::is_integral<T>::value) {
76+
size = static_cast<int64>(
77+
(std::abs(limit - start) + std::abs(delta) - 1) / std::abs(delta));
78+
} else {
79+
size = static_cast<int64>(std::ceil(std::abs((limit - start) / delta)));
80+
}
7881
Tensor* out = nullptr;
7982
OP_REQUIRES_OK(context,
8083
context->allocate_output(0, TensorShape({size}), &out));

Diff for: tensorflow/python/kernel_tests/init_ops_test.py

+8
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
from tensorflow.python.framework import constant_op
2525
from tensorflow.python.framework import dtypes
26+
from tensorflow.python.framework import errors_impl
2627
from tensorflow.python.framework import ops
2728
from tensorflow.python.framework import random_seed
2829
from tensorflow.python.framework import test_util
@@ -542,6 +543,13 @@ def testMixedDType(self):
542543
constant_op.constant(4, dtype=dtypes.int32), dtype=dtypes.int64)
543544
self.assertAllEqual(self.evaluate(tf_ans), np.array([0, 1, 2, 3]))
544545

546+
def testLargeLimits(self):
547+
# Test case for GitHub issue 46913.
548+
with self.session():
549+
with self.assertRaises(errors_impl.ResourceExhaustedError):
550+
v = math_ops.range(0, 9223372036854775807)
551+
self.evaluate(v)
552+
545553

546554
# TODO(vrv): move to sequence_ops_test?
547555
class LinSpaceTest(test.TestCase):

0 commit comments

Comments
 (0)