36 lines
1.8 KiB
Diff
36 lines
1.8 KiB
Diff
|
From f2a673bd34f0d64b8e40a551ac78989d16daad09 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Fri, 30 Jul 2021 19:00:00 -0700
|
||
|
|
Subject: [PATCH] Add missing validation to `matrix_diag_op.cc`
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 387923533
|
||
|
|
Change-Id: Idfffeb328d5f9c6748d992d28a56d6e9e45103a0
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/matrix_diag_op.cc | 6 ++++++
|
||
|
|
1 file changed, 6 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/matrix_diag_op.cc b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
index 2fd361e715b9b..b61dbe96a5d6a 100644
|
||
|
|
--- a/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
@@ -73,6 +73,9 @@ class MatrixDiagPartOp : public OpKernel {
|
||
|
|
errors::InvalidArgument(
|
||
|
|
"diag_index must be a scalar or vector, received shape: ",
|
||
|
|
diag_index.shape().DebugString()));
|
||
|
|
+ OP_REQUIRES(context, diag_index.NumElements() > 0,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "Expected diag_index to have at least 1 element"));
|
||
|
|
lower_diag_index = diag_index.flat<int32>()(0);
|
||
|
|
upper_diag_index = lower_diag_index;
|
||
|
|
if (TensorShapeUtils::IsVector(diag_index.shape())) {
|
||
|
|
@@ -179,6 +182,9 @@ class MatrixDiagOp : public OpKernel {
|
||
|
|
errors::InvalidArgument(
|
||
|
|
"diag_index must be a scalar or vector, received shape: ",
|
||
|
|
diag_index.shape().DebugString()));
|
||
|
|
+ OP_REQUIRES(context, diag_index.NumElements() > 0,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "Expected diag_index to have at least 1 element"));
|
||
|
|
lower_diag_index = diag_index.flat<int32>()(0);
|
||
|
|
upper_diag_index = lower_diag_index;
|
||
|
|
if (TensorShapeUtils::IsVector(diag_index.shape())) {
|