Test for case that would lead to undefined behaviour and throw error …

…if found
tensorflow · Nov 2, 2021 · 8a363d6 · 8a363d6
1 parent 2a967b1
commit 8a363d6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
@@ -71,13 +71,18 @@ class RangeOp : public OpKernel {
           errors::InvalidArgument(
               "Requires start >= limit when delta < 0: ", start, "/", limit));
     }
-    int64_t size = 0;
-    if (std::is_integral<T>::value) {
-      size = static_cast<int64_t>(
-          (std::abs(limit - start) + std::abs(delta) - 1) / std::abs(delta));
-    } else {
-      size = static_cast<int64_t>(std::ceil(std::abs((limit - start) / delta)));
-    }
+    auto size_auto = (std::is_integral<T>::value
+	    ? (Eigen::numext::abs(limit - start) + 
+		    Eigen::numext::abs(delta) - T(1)) / 
+		    Eigen::numext::abs(delta)
+	    : Eigen::numext::ceil(
+		    Eigen::numext::abs((limit - start) / delta)));
+    OP_REQUIRES(context, size_auto <= std::numeric_limits<int64_t>::max(),
+          errors::InvalidArgument("Requires ((limit - start) / delta) <= ",
+				  std::numeric_limits<int64_t>::max()));
+
+    int64_t size = static_cast<int64_t>(size_auto);
+
     TensorShape shape;
     OP_REQUIRES_OK(context, shape.AddDimWithStatus(size));
     Tensor* out = nullptr;

diff --git a/tensorflow/core/ops/math_ops.cc b/tensorflow/core/ops/math_ops.cc
@@ -1481,6 +1481,13 @@ Status RangeSize(const Tensor* start_t, const Tensor* limit_t,
                       Eigen::numext::abs(delta))
                    : (Eigen::numext::ceil(
                          Eigen::numext::abs((limit - start) / delta))));
+
+  // Undefined behaviour if size will not fit into int64_t
+  if (size > std::numeric_limits<int64_t>::max()) {
+    return errors::InvalidArgument(
+        "Requires ((limit - start) / delta) <= ", std::numeric_limits<int64_t>::max());
+  }
+
   c->set_output(0, c->Vector(static_cast<int64_t>(size)));
   return Status::OK();
 }

diff --git a/tensorflow/python/kernel_tests/array_ops/init_ops_test.py b/tensorflow/python/kernel_tests/array_ops/init_ops_test.py
@@ -549,7 +549,7 @@ def testLargeLimits(self):
   def testLargeStarts(self):
     # Test case for GitHub issue 46899.
     with self.session():
-      with self.assertRaises(errors_impl.InternalError):
+      with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
         v = math_ops.range(start=-1e+38, limit=1)
         self.evaluate(v)