Ensure non-empty padding_value input to tf.raw_ops.MatrixDiagPartV2, if a padding_value is input

pak-laura · tensorflower-gardener · commit 482da92095c4 · 2021-08-02T15:21:23.000-07:00
PiperOrigin-RevId: 388314614
Change-Id: If0b51ad58d5d8543a6be6ce8f42ae4755c80d55f
diff --git a/tensorflow/core/kernels/linalg/matrix_diag_op.cc b/tensorflow/core/kernels/linalg/matrix_diag_op.cc
@@ -89,7 +89,10 @@ class MatrixDiagPartOp : public OpKernel {
           upper_diag_index = diag_index.flat<int32>()(1);
         }
       }
-      padding_value = context->input(2).flat<T>()(0);
+      const Tensor& padding_in = context->input(2);
+      OP_REQUIRES(context, padding_in.NumElements() == 1,
+                  errors::InvalidArgument("Padding must be scalar."));
+      padding_value = padding_in.flat<T>()(0);
     }
     const TensorShape& input_shape = input.shape();
 

Original file line number	Diff line number	Diff line change
`@@ -89,7 +89,10 @@ class MatrixDiagPartOp : public OpKernel {`
`89`	`89`	`upper_diag_index = diag_index.flat<int32>()(1);`
`90`	`90`	`}`
`91`	`91`	`}`
`92`		`- padding_value = context->input(2).flat<T>()(0);`
	`92`	`+ const Tensor& padding_in = context->input(2);`
	`93`	`+ OP_REQUIRES(context, padding_in.NumElements() == 1,`
	`94`	`+ errors::InvalidArgument("Padding must be scalar."));`
	`95`	`+ padding_value = padding_in.flat<T>()(0);`
`93`	`96`	`}`
`94`	`97`	`const TensorShape& input_shape = input.shape();`
`95`	`98`