66 lines
2.7 KiB
Diff
66 lines
2.7 KiB
Diff
|
From 3f6fe4dfef6f57e768260b48166c27d148f3015f Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Tue, 4 May 2021 18:33:28 -0700
|
||
|
|
Subject: [PATCH] Add missing validations in dillation ops.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 372037158
|
||
|
|
Change-Id: I4ee304c84a02550c030288a6534000b934fc1599
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/dilation_ops.cc | 15 +++++++++++----
|
||
|
|
1 file changed, 11 insertions(+), 4 deletions(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/dilation_ops.cc b/tensorflow/core/kernels/dilation_ops.cc
|
||
|
|
index 738ea31d555d5..996ddb62bfefe 100644
|
||
|
|
--- a/tensorflow/core/kernels/dilation_ops.cc
|
||
|
|
+++ b/tensorflow/core/kernels/dilation_ops.cc
|
||
|
|
@@ -130,6 +130,7 @@ class DilationOp : public OpKernel {
|
||
|
|
ParseSizes(context, strides_, rates_, padding_, &stride_rows, &stride_cols,
|
||
|
|
&rate_rows, &rate_cols, &pad_top, &pad_left, &out_rows,
|
||
|
|
&out_cols);
|
||
|
|
+ if (!context->status().ok()) return;
|
||
|
|
|
||
|
|
// Output tensor is of the following dimensions:
|
||
|
|
// [ batch, out_rows, out_cols, depth ]
|
||
|
|
@@ -229,6 +230,7 @@ class DilationBackpropInputOp : public OpKernel {
|
||
|
|
ParseSizes(context, strides_, rates_, padding_, &stride_rows, &stride_cols,
|
||
|
|
&rate_rows, &rate_cols, &pad_top, &pad_left, &out_rows,
|
||
|
|
&out_cols);
|
||
|
|
+ if (!context->status().ok()) return;
|
||
|
|
|
||
|
|
// Verify that the incoming gradient tensor has the expected size
|
||
|
|
// [ batch, out_rows, out_cols, depth ]
|
||
|
|
@@ -318,8 +320,10 @@ struct DilationBackpropInput<CPUDevice, T> {
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
- in_backprop(b, h_in_max, w_in_max, d) +=
|
||
|
|
- out_backprop(b, h_out, w_out, d);
|
||
|
|
+ if (h_in_max < input_rows && w_in_max < input_cols) {
|
||
|
|
+ in_backprop(b, h_in_max, w_in_max, d) +=
|
||
|
|
+ out_backprop(b, h_out, w_out, d);
|
||
|
|
+ }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
@@ -349,6 +353,7 @@ class DilationBackpropFilterOp : public OpKernel {
|
||
|
|
ParseSizes(context, strides_, rates_, padding_, &stride_rows, &stride_cols,
|
||
|
|
&rate_rows, &rate_cols, &pad_top, &pad_left, &out_rows,
|
||
|
|
&out_cols);
|
||
|
|
+ if (!context->status().ok()) return;
|
||
|
|
|
||
|
|
// Verify that the incoming gradient tensor has the expected size
|
||
|
|
// [ batch, out_rows, out_cols, depth ]
|
||
|
|
@@ -438,8 +443,10 @@ struct DilationBackpropFilter<CPUDevice, T> {
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
- filter_backprop(h_max, w_max, d) +=
|
||
|
|
- out_backprop(b, h_out, w_out, d);
|
||
|
|
+ if (h_max < filter_rows && w_max < filter_cols) {
|
||
|
|
+ filter_backprop(h_max, w_max, d) +=
|
||
|
|
+ out_backprop(b, h_out, w_out, d);
|
||
|
|
+ }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|