Merge pull request #49544 from geetachavan1/cherrypicks_UUYUB

Fix FPE issue in external Eigen source code issue with `tf.raw_ops.SparseMatMul`.
tensorflow · May 27, 2021 · a525ac6 · a525ac6
2 parents 99931bc + dd63731
commit a525ac6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/sparse_matmul_op.cc b/tensorflow/core/kernels/sparse_matmul_op.cc
@@ -1039,6 +1039,10 @@ class SparseMatMulOp : public OpKernel {
     if (transpose_b) {
       // TODO(agarwal): avoid transposing the matrix here and directly handle
       // transpose in CreateDenseSlices.
+      OP_REQUIRES(ctx, right->dim_size(0) != 0,
+                  errors::InvalidArgument("b has an entry 0 in it's shape."));
+      OP_REQUIRES(ctx, right->dim_size(1) != 0,
+                  errors::InvalidArgument("b has an entry 0 in it's shape."));
       right_tr.reset(
           new Tensor(right->dtype(),
                      TensorShape({right->dim_size(1), right->dim_size(0)})));