From 7ae2af34087fb4b5c8915279efd03da3b81028bc Mon Sep 17 00:00:00 2001 From: Amit Patankar 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(); const auto shape_vec = shape_t->vec();