From b0375e267c55ae97e65b50ce51200e012ac3e8ad Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Sat, 6 Jun 2020 01:24:20 +0000
Subject: [PATCH 1/2] Fix floating point exception with tf.unravel_index
This PR tries to address the issue raised in 40204 where
`tf.unravel_index` caused floating point exception when any one
element is 0 in `dims`.
The issue is that `indices` in `tf.unravel_index` should make sure
it is not out of boundary compared to `dims`.
This PR fixes the issue by adding a check before hand, though Eigen.
This PR fixes 40204.
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
tensorflow/core/kernels/unravel_index_op.cc | 18 ++++++++++++++++++
.../python/kernel_tests/array_ops_test.py | 10 ++++++++++
2 files changed, 28 insertions(+)
diff --git a/tensorflow/core/kernels/unravel_index_op.cc b/tensorflow/core/kernels/unravel_index_op.cc
index 8d839ba85a7..d41915e5c14 100644
--- a/tensorflow/core/kernels/unravel_index_op.cc
+++ b/tensorflow/core/kernels/unravel_index_op.cc
@@ -54,6 +54,24 @@ class UnravelIndexOp : public OpKernel {
auto dims = dims_tensor.vec<Tidx>();
+ // Chek to make sure indices is not out of boundary
+ Eigen::Tensor<bool, 0, Eigen::RowMajor> check;
+ if (TensorShapeUtils::IsScalar(indices_tensor.shape())) {
+ auto indices = indices_tensor.scalar<Tidx>();
+ auto dims_prod = dims.prod();
+ check = (indices < dims_prod).all();
+ } else {
+ auto indices = indices_tensor.vec<Tidx>();
+ auto dims_prod = dims.prod()
+ .reshape(Eigen::array<Eigen::Index, 1>({1}))
+ .broadcast(
+ Eigen::array<Eigen::Index, 1>({indices_tensor.NumElements()}));
+ check = (indices < dims_prod).all();
+ }
+ OP_REQUIRES(
+ ctx, check(),
+ errors::InvalidArgument("index is out of bound as with dims"));
+
Eigen::array<bool, 1> reverse({true});
Tensor strides_tensor;
diff --git a/tensorflow/python/kernel_tests/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops_test.py
index ec3ed932996..caf05042557 100644
--- a/tensorflow/python/kernel_tests/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops_test.py
@@ -1384,6 +1384,16 @@ class UnravelIndexTest(test_util.TensorFlowTestCase):
out_3 = array_ops.unravel_index(indices_3, dims_3)
self.assertAllEqual(out_3.eval(), [[3, 6, 6], [4, 5, 1]])
+ # Test case for GitHub issue 40204.
+ def testUnravelIndexZeroDim(self):
+ with self.cached_session():
+ for dtype in [dtypes.int32, dtypes.int64]:
+ with self.assertRaisesRegexp(
+ errors.InvalidArgumentError, "index is out of bound as with dims"):
+ indices = constant_op.constant([2, 5, 7], dtype=dtype)
+ dims = constant_op.constant([3, 0], dtype=dtype)
+ self.evaluate(array_ops.unravel_index(indices=indices, dims=dims))
+
class GuaranteeConstOpTest(test_util.TensorFlowTestCase):
From 43930312404d885f1d6de086120712ddb633afdb Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Sun, 7 Jun 2020 17:20:27 +0000
Subject: [PATCH 2/2] Update the boundary check in unravel_index to use
std::all_of instead of eigen
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
tensorflow/core/kernels/unravel_index_op.cc | 22 ++++++++-------------
1 file changed, 8 insertions(+), 14 deletions(-)
diff --git a/tensorflow/core/kernels/unravel_index_op.cc b/tensorflow/core/kernels/unravel_index_op.cc
index d41915e5c14..fb8d6703d57 100644
--- a/tensorflow/core/kernels/unravel_index_op.cc
+++ b/tensorflow/core/kernels/unravel_index_op.cc
@@ -55,21 +55,15 @@ class UnravelIndexOp : public OpKernel {
auto dims = dims_tensor.vec<Tidx>();
// Chek to make sure indices is not out of boundary
- Eigen::Tensor<bool, 0, Eigen::RowMajor> check;
- if (TensorShapeUtils::IsScalar(indices_tensor.shape())) {
- auto indices = indices_tensor.scalar<Tidx>();
- auto dims_prod = dims.prod();
- check = (indices < dims_prod).all();
- } else {
- auto indices = indices_tensor.vec<Tidx>();
- auto dims_prod = dims.prod()
- .reshape(Eigen::array<Eigen::Index, 1>({1}))
- .broadcast(
- Eigen::array<Eigen::Index, 1>({indices_tensor.NumElements()}));
- check = (indices < dims_prod).all();
- }
+ Eigen::Tensor<Tidx, 0, Eigen::RowMajor> dims_prod_eigen = dims.prod();
+ Tidx dims_prod = dims_prod_eigen();
+ const Tidx* indices = indices_tensor.flat<Tidx>().data();
+ int64 size = indices_tensor.NumElements();
+ bool check = std::all_of(indices, indices + size, [&](Tidx index) {
+ return index < dims_prod;
+ });
OP_REQUIRES(
- ctx, check(),
+ ctx, check,
errors::InvalidArgument("index is out of bound as with dims"));
Eigen::array<bool, 1> reverse({true});