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>();