42 lines
1.8 KiB
Diff
42 lines
1.8 KiB
Diff
|
From a7116dd3913c4a4afd2a3a938573aa7c785fdfc6 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Sat, 17 Apr 2021 20:55:53 -0700
|
||
|
|
Subject: [PATCH] Validate `MatrixDiagV{2,3}` arguments to prevent breakage.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 369056033
|
||
|
|
Change-Id: Ic2018c297d3dd6f252dc1dd3667f1ed5cb1eaa42
|
||
|
|
---
|
||
|
|
.../core/kernels/matrix_diag_op.cc | 19 ++++++++++++++++---
|
||
|
|
1 file changed, 16 insertions(+), 3 deletions(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/matrix_diag_op.cc b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
index 69cc8170793ae..d4eb589836a85 100644
|
||
|
|
--- a/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
@@ -192,9 +192,22 @@ class MatrixDiagOp : public OpKernel {
|
||
|
|
upper_diag_index = diag_index.flat<int32>()(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
- num_rows = context->input(2).flat<int32>()(0);
|
||
|
|
- num_cols = context->input(3).flat<int32>()(0);
|
||
|
|
- padding_value = context->input(4).flat<T>()(0);
|
||
|
|
+
|
||
|
|
+ auto& num_rows_tensor = context->input(2);
|
||
|
|
+ OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_rows_tensor.shape()),
|
||
|
|
+ errors::InvalidArgument("num_rows must be a scalar"));
|
||
|
|
+ num_rows = num_rows_tensor.flat<int32>()(0);
|
||
|
|
+
|
||
|
|
+ auto& num_cols_tensor = context->input(3);
|
||
|
|
+ OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_cols_tensor.shape()),
|
||
|
|
+ errors::InvalidArgument("num_cols must be a scalar"));
|
||
|
|
+ num_cols = num_cols_tensor.flat<int32>()(0);
|
||
|
|
+
|
||
|
|
+ auto& padding_value_tensor = context->input(4);
|
||
|
|
+ OP_REQUIRES(context,
|
||
|
|
+ TensorShapeUtils::IsScalar(padding_value_tensor.shape()),
|
||
|
|
+ errors::InvalidArgument("padding_value must be a scalar"));
|
||
|
|
+ padding_value = padding_value_tensor.flat<T>()(0);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Size validations.
|