Merge pull request #52707 from elfringham:init_ops_test_fix

PiperOrigin-RevId: 416941851 Change-Id: Iefa5a9b841b053b36f6b105cd82c9d32d5e47850
tensorflow · Jan 25, 2022 · e896767 · e896767
1 parent e3adc2d
commit e896767
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
@@ -71,13 +71,19 @@ class RangeOp : public OpKernel {
           errors::InvalidArgument(
               "Requires start >= limit when delta < 0: ", start, "/", limit));
     }
-    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)));
-    }
+    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>::max(),
+        errors::InvalidArgument("Requires ((limit - start) / delta) <= ",
+                                std::numeric_limits<int64>::max()));
+
+    int64 size = static_cast<int64>(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>(size)));
   return Status::OK();
 }

diff --git a/tensorflow/python/kernel_tests/init_ops_test.py b/tensorflow/python/kernel_tests/init_ops_test.py
@@ -546,7 +546,7 @@ def testMixedDType(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)