Skip to content
Permalink
Browse files Browse the repository at this point in the history
Validate that matrix dimension sizes in SparseMatMul are positive.
PiperOrigin-RevId: 401149683
Change-Id: Ib33eafc561a39c8741ece80b2edce6d4aae9a57d
  • Loading branch information
penpornk authored and tensorflower-gardener committed Oct 6, 2021
1 parent d4fdd78 commit e6cf28c
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions tensorflow/core/kernels/sparse_matmul_op.cc
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit e6cf28c

Please sign in to comment.