29 lines
1.1 KiB
Diff
29 lines
1.1 KiB
Diff
|
From 482da92095c4d48f8784b1f00dda4f81c28d2988 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Laura Pak <lpak@google.com>
|
||
|
|
Date: Mon, 2 Aug 2021 15:07:31 -0700
|
||
|
|
Subject: [PATCH] Ensure non-empty padding_value input to
|
||
|
|
tf.raw_ops.MatrixDiagPartV2, if a padding_value is input
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 388314614
|
||
|
|
Change-Id: If0b51ad58d5d8543a6be6ce8f42ae4755c80d55f
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/matrix_diag_op.cc | 5 ++++-
|
||
|
|
1 file changed, 4 insertions(+), 1 deletion(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/matrix_diag_op.cc b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
index b61dbe96a5d6a..1506df1712138 100644
|
||
|
|
--- a/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/matrix_diag_op.cc
|
||
|
|
@@ -86,7 +86,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();
|
||
|
|
|