Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix overflow issue
We're also decreasing the test numbers by 1 as the overflow threshold is 1 less.

PiperOrigin-RevId: 449856148
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 19, 2022
1 parent 52970bd commit 37e6453
Showing 1 changed file with 12 additions and 11 deletions.
23 changes: 12 additions & 11 deletions tensorflow/core/ops/math_ops.cc
Expand Up @@ -1487,17 +1487,18 @@ Status RangeSize(const Tensor* start_t, const Tensor* limit_t,
return errors::InvalidArgument("Requires delta != 0");
}

auto size = (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))));

// 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());
int64_t size;
if (std::is_integral<T>::value) {
size = Eigen::divup(static_cast<int64_t>(Eigen::numext::abs(limit - start)),
static_cast<int64_t>(Eigen::numext::abs(delta)));
} else {
auto size_auto =
Eigen::numext::ceil(Eigen::numext::abs((limit - start) / delta));
if (size_auto > std::numeric_limits<int64_t>::max()) {
return errors::InvalidArgument("Requires ((limit - start) / delta) <= ",
std::numeric_limits<int64_t>::max());
}
size = static_cast<int64_t>(size_auto);
}

c->set_output(0, c->Vector(static_cast<int64_t>(size)));
Expand Down

0 comments on commit 37e6453

Please sign in to comment.