36 lines
1.4 KiB
Diff
36 lines
1.4 KiB
Diff
From 1e922ccdf6bf46a3a52641f99fd47d54c1decd13 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
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);
|
|
}
|