Merge pull request #49630 from geetachavan1/cherrypicks_IXSZF

Fix heap-buffer-overflow issue with `tf.raw_ops.SparseDenseCwiseMul`.
tensorflow · May 29, 2021 · cf72480 · cf72480
2 parents b97dc07 + 971743b
commit cf72480
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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>();