Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap-buffer-overflow issue with tf.raw_ops.SparseDenseCwiseMul.
PiperOrigin-RevId: 372054410
Change-Id: Ifcce0491e2e3816838c87e73be30a1e61b65174d
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed May 5, 2021
1 parent 38178a2 commit 7ae2af3
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
Expand Up @@ -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>();
Expand Down

0 comments on commit 7ae2af3

Please sign in to comment.