44 lines
1.9 KiB
Diff
44 lines
1.9 KiB
Diff
|
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.
|