From e6cf28c72ba2eb949ca950d834dd6d66bb01cfae Mon Sep 17 00:00:00 2001
From: Penporn Koanantakool <penporn@google.com>
Date: Tue, 5 Oct 2021 21:54:15 -0700
Subject: [PATCH] Validate that matrix dimension sizes in SparseMatMul are
 positive.

PiperOrigin-RevId: 401149683
Change-Id: Ib33eafc561a39c8741ece80b2edce6d4aae9a57d
---
 tensorflow/core/kernels/sparse_matmul_op.cc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/core/kernels/sparse_matmul_op.cc b/tensorflow/core/kernels/sparse_matmul_op.cc
index a02afafa33e3a..6bf9dfa3d8bb7 100644
--- a/tensorflow/core/kernels/sparse_matmul_op.cc
+++ b/tensorflow/core/kernels/sparse_matmul_op.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/fill_functor.h"
 #include "tensorflow/core/lib/core/blocking_counter.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/mutex.h"
@@ -980,9 +981,18 @@ class SparseMatMulOp : public OpKernel {
                 errors::InvalidArgument(
                     "Matrix size incompatible: a: ", a.shape().DebugString(),
                     ", b: ", b.shape().DebugString()));
+    OP_REQUIRES(ctx, m >= 0 && n >= 0 && k >= 0,
+                errors::InvalidArgument(
+                    "Matrix dimensions cannot be negative: a: ",
+                    a.shape().DebugString(), ", b: ", b.shape().DebugString()));
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({m, n}), &output));
 
+    // Return early if at least one of the output dimension size is 0.
+    if (m == 0 || n == 0) {
+      return;
+    }
+
     if (k == 0) {
       // If the inner dimension k in the matrix multiplication is zero, we fill
       // the output with zeros.