From 0666d8bb711b41c9f03dec238d7d165bc946fc70 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Mon, 14 Dec 2020 23:11:16 +0000 Subject: [PATCH 1/3] Prevent crash of tensorflow if shape is too large for tf.sparse.reorder This PR tries to address the issue raised in 45392 where tensorflow crashes if shape of sparse tensor is too large for tf.sparse.reorder This PR adds additional checks and exit gracefully if the shape is too large. This PR fixes 45392. Signed-off-by: Yong Tang --- tensorflow/core/kernels/sparse_reorder_op.cc | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tensorflow/core/kernels/sparse_reorder_op.cc b/tensorflow/core/kernels/sparse_reorder_op.cc index 6f9065827fd..465bff8d3cf 100644 --- a/tensorflow/core/kernels/sparse_reorder_op.cc +++ b/tensorflow/core/kernels/sparse_reorder_op.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/util/sparse/sparse_tensor.h" +#include "tensorflow/core/util/overflow.h" namespace tensorflow { @@ -54,6 +55,18 @@ class SparseReorderOp : public OpKernel { "Input shape should be a vector but received shape ", input_shape_in.shape().DebugString())); + // Check if the sparse tensor input shape is valid + int64 total = 1; + for (int64 i = 0; i < input_shape_in.NumElements(); ++i) { + int dim = input_shape_in.vec()(i); + OP_REQUIRES(context, (dim >= 0), + errors::InvalidArgument("Dimension ", dim, " must be >= 0")); + total = MultiplyWithoutOverflow(total, dim); + OP_REQUIRES(context, (total > 0), + errors::InvalidArgument( + "Shape would have more than 2**63 - 1 elements")); + } + const TensorShape input_shape(input_shape_in.vec()); gtl::InlinedVector std_order(input_shape.dims()); From f137ea2a23ac1e171224a16d59efc060ea06a616 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Mon, 14 Dec 2020 23:12:56 +0000 Subject: [PATCH 2/3] Add test case for GitHub issue 45392. Signed-off-by: Yong Tang --- .../python/kernel_tests/sparse_reorder_op_test.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorflow/python/kernel_tests/sparse_reorder_op_test.py b/tensorflow/python/kernel_tests/sparse_reorder_op_test.py index 93fcc6a18e6..d05fd4e9985 100644 --- a/tensorflow/python/kernel_tests/sparse_reorder_op_test.py +++ b/tensorflow/python/kernel_tests/sparse_reorder_op_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import numpy as np from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors_impl from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops @@ -124,6 +125,17 @@ class SparseReorderTest(test.TestCase): x_init_value=input_val.values) self.assertLess(err, 1e-11) + def testShapeOverflow(self): + # Test case for GitHub issue 45392 + sp_input = sparse_tensor.SparseTensor( + indices=[[0, 0, 0, 0, 0, 0]], values=[0.0], + dense_shape=[4096, 4096, 4096, 4096, 4096, 4096]) + self.assertAllEqual( + (4096, 4096, 4096, 4096, 4096, 4096), sp_input.get_shape()) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "Shape would have more than"): + sp_output = sparse_ops.sparse_reorder(sp_input) + if __name__ == "__main__": test.main() From 63690f568c87237a76d91af21177456fcd8fdfd1 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Thu, 17 Dec 2020 00:00:10 +0000 Subject: [PATCH 3/3] Allow sparse tensor with reorder work on potentially large dims, Update: review comments addressed. Signed-off-by: Yong Tang --- tensorflow/core/kernels/sparse_reorder_op.cc | 18 +++--------------- .../kernel_tests/sparse_reorder_op_test.py | 7 +++---- 2 files changed, 6 insertions(+), 19 deletions(-) diff --git a/tensorflow/core/kernels/sparse_reorder_op.cc b/tensorflow/core/kernels/sparse_reorder_op.cc index 465bff8d3cf..b7fb56185c3 100644 --- a/tensorflow/core/kernels/sparse_reorder_op.cc +++ b/tensorflow/core/kernels/sparse_reorder_op.cc @@ -27,7 +27,6 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/util/sparse/sparse_tensor.h" -#include "tensorflow/core/util/overflow.h" namespace tensorflow { @@ -55,21 +54,10 @@ class SparseReorderOp : public OpKernel { "Input shape should be a vector but received shape ", input_shape_in.shape().DebugString())); - // Check if the sparse tensor input shape is valid - int64 total = 1; - for (int64 i = 0; i < input_shape_in.NumElements(); ++i) { - int dim = input_shape_in.vec()(i); - OP_REQUIRES(context, (dim >= 0), - errors::InvalidArgument("Dimension ", dim, " must be >= 0")); - total = MultiplyWithoutOverflow(total, dim); - OP_REQUIRES(context, (total > 0), - errors::InvalidArgument( - "Shape would have more than 2**63 - 1 elements")); - } + gtl::ArraySlice input_shape( + input_shape_in.vec().data(), input_shape_in.NumElements()); - const TensorShape input_shape(input_shape_in.vec()); - - gtl::InlinedVector std_order(input_shape.dims()); + gtl::InlinedVector std_order(input_shape.size()); std::iota(std_order.begin(), std_order.end(), 0); // Check if the sparse tensor is already ordered correctly diff --git a/tensorflow/python/kernel_tests/sparse_reorder_op_test.py b/tensorflow/python/kernel_tests/sparse_reorder_op_test.py index d05fd4e9985..c889fb14fbd 100644 --- a/tensorflow/python/kernel_tests/sparse_reorder_op_test.py +++ b/tensorflow/python/kernel_tests/sparse_reorder_op_test.py @@ -21,7 +21,6 @@ from __future__ import print_function import numpy as np from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors_impl from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops @@ -132,9 +131,9 @@ class SparseReorderTest(test.TestCase): dense_shape=[4096, 4096, 4096, 4096, 4096, 4096]) self.assertAllEqual( (4096, 4096, 4096, 4096, 4096, 4096), sp_input.get_shape()) - with self.assertRaisesRegex(errors_impl.InvalidArgumentError, - "Shape would have more than"): - sp_output = sparse_ops.sparse_reorder(sp_input) + sp_output = sparse_ops.sparse_reorder(sp_input) + self.assertAllEqual( + (4096, 4096, 4096, 4096, 4096, 4096), sp_output.get_shape()) if __name__ == "__main__":