Merge pull request #52707 from elfringham:init_ops_test_fix

tensorflower-gardener · tensorflower-gardener · commit f0147751fd5d · 2021-12-16T19:26:39.000-08:00
PiperOrigin-RevId: 416941851
Change-Id: Iefa5a9b841b053b36f6b105cd82c9d32d5e47850
diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
@@ -91,13 +91,19 @@ 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
@@ -1489,6 +1489,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
@@ -548,7 +548,7 @@ def testLargeLimits(self):
   def testLargeStarts(self):
     # Test case for GitHub issue 46899.
     with self.session():
-      with self.assertRaises(errors_impl.InvalidArgumentError):
+      with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
         v = math_ops.range(start=-1e+38, limit=1)
         self.evaluate(v)
 
