29 lines
1.3 KiB
Diff
29 lines
1.3 KiB
Diff
|
From 7ae2af34087fb4b5c8915279efd03da3b81028bc Mon Sep 17 00:00:00 2001
|
||
|
|
From: Amit Patankar <amitpatankar@google.com>
|
||
|
|
Date: Tue, 4 May 2021 21:30:50 -0700
|
||
|
|
Subject: [PATCH] Fix heap-buffer-overflow issue with
|
||
|
|
`tf.raw_ops.SparseDenseCwiseMul`.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 372054410
|
||
|
|
Change-Id: Ifcce0491e2e3816838c87e73be30a1e61b65174d
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/sparse_dense_binary_op_shared.cc | 5 +++++
|
||
|
|
1 file changed, 5 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc b/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
|
||
|
|
index 3a5e66a0e73ea..dac4a3d3e6bfc 100644
|
||
|
|
--- a/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
|
||
|
|
+++ b/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
|
||
|
|
@@ -78,6 +78,11 @@ class SparseDenseBinaryOpShared : public OpKernel {
|
||
|
|
"but received shapes: ",
|
||
|
|
values_t->shape().DebugString(), " and ",
|
||
|
|
shape_t->shape().DebugString()));
|
||
|
|
+ OP_REQUIRES(
|
||
|
|
+ ctx, values_t->dim_size(0) == indices_t->dim_size(0),
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "The first dimension of values and indices should match. (",
|
||
|
|
+ values_t->dim_size(0), " vs. ", indices_t->dim_size(0), ")"));
|
||
|
|
|
||
|
|
const auto indices_mat = indices_t->matrix<int64>();
|
||
|
|
const auto shape_vec = shape_t->vec<int64>();
|