commit 3f6fe4dfef6f57e768260b48166c27d148f3015f
author Mihai Maruseac <mihaimaruseac@google.com> Tue May 04 18:33:28 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Tue May 04 18:41:15 2021 -0700
tree 9519dc67f60604181bd17a549bfb32c366b1d9c8
parent afd954e65f15aea4d438d0a219136fc4a63a573d

Add missing validations in dillation ops.

PiperOrigin-RevId: 372037158
Change-Id: I4ee304c84a02550c030288a6534000b934fc1599

tensorflow/core/kernels/dilation_ops.cc

diff --git a/tensorflow/core/kernels/dilation_ops.cc b/tensorflow/core/kernels/dilation_ops.cc
index 738ea31..996ddb6 100644
--- a/tensorflow/core/kernels/dilation_ops.cc
+++ b/tensorflow/core/kernels/dilation_ops.cc
@@ -130,6 +130,7 @@
     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 @@
     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 @@
                 }
               }
             }
-            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 @@
     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 @@
                 }
               }
             }
-            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);
+            }
           }
         }
       }