Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add missing validations in dillation ops.
PiperOrigin-RevId: 372037158
Change-Id: I4ee304c84a02550c030288a6534000b934fc1599
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 5, 2021
1 parent afd954e commit 3f6fe4d
Showing 1 changed file with 11 additions and 4 deletions.
15 changes: 11 additions & 4 deletions tensorflow/core/kernels/dilation_ops.cc
Expand Up @@ -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 ]
Expand Down Expand Up @@ -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 ]
Expand Down Expand Up @@ -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);
}
}
}
}
Expand Down Expand Up @@ -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 ]
Expand Down Expand Up @@ -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);
}
}
}
}
Expand Down

0 comments on commit 3f6fe4d

Please sign in to comment.