From 1e922ccdf6bf46a3a52641f99fd47d54c1decd13 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Fri, 23 Apr 2021 10:41:12 -0700 Subject: [PATCH] Fix crash in `SparseTensorToCSRSparseMatrixCPUFunctor` PiperOrigin-RevId: 370110290 Change-Id: I4451e92661a55c2180f80d38b67a9b50bf5edec5 --- tensorflow/core/kernels/sparse/kernels.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/kernels/sparse/kernels.cc b/tensorflow/core/kernels/sparse/kernels.cc index 0eea9f1feed5c..dff9aeb83ccfe 100644 --- a/tensorflow/core/kernels/sparse/kernels.cc +++ b/tensorflow/core/kernels/sparse/kernels.cc @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/errors.h" namespace tensorflow { namespace functor { @@ -63,6 +64,11 @@ Status SparseTensorToCSRSparseMatrixCPUFunctor::operator()( for (int64 i = 0; i < total_nnz; ++i) { // For now, the rows pointers store the corresponding row counts. + int64 ix = indices(i, 0) + 1; + if (ix >= csr_row_ptr.size()) { + return errors::InvalidArgument("Got an index ", ix, + " that is outside of csr_row_ptr"); + } csr_row_ptr(indices(i, 0) + 1) += 1; csr_col_ind(i) = indices(i, 1); }