| 190 | } |
| 191 | |
| 192 | XlaOp AvgPoolGrad(XlaOp out_backprop, absl::Span<const int64> gradients_size, |
| 193 | absl::Span<const int64> kernel_size, |
| 194 | absl::Span<const int64> stride, |
| 195 | absl::Span<const std::pair<int64, int64>> spatial_padding, |
| 196 | const TensorFormat& data_format, |
| 197 | const bool counts_include_padding) { |
| 198 | XlaBuilder* b = out_backprop.builder(); |
| 199 | return b->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 200 | const int num_dims = kernel_size.size(); |
| 201 | |
| 202 | if (gradients_size.size() != num_dims) { |
| 203 | return tensorflow::errors::InvalidArgument("gradients must be ", num_dims, |
| 204 | "-dimensional"); |
| 205 | } |
| 206 | |
| 207 | TF_ASSIGN_OR_RETURN(Shape out_backprop_xla_shape, |
| 208 | b->GetShape(out_backprop)); |
| 209 | if (out_backprop_xla_shape.dimensions().size() != num_dims) { |
| 210 | return tensorflow::errors::InvalidArgument("out_backprop must be ", |
| 211 | num_dims, "-dimensional"); |
| 212 | } |
| 213 | |
| 214 | // We can think of average-pooling as: |
| 215 | // * a convolution with a kernel consisting entirely of 1s, where the |
| 216 | // input feature and output feature are equal, and 0s everywhere else. |
| 217 | // * followed by dividing by the counts. |
| 218 | // |
| 219 | // This then gives us an algorithm to build the gradient: |
| 220 | // * divide out_backprop by the counts, followed by |
| 221 | // * Conv2DBackpropInput specialized for that kernel, which simplifies to |
| 222 | // a Pad and a ReduceWindow. |
| 223 | // |
| 224 | // For an explanation of backpropagation for convolution, see the comments |
| 225 | // in third_party/tensorflow/core/kernels/conv_grad_ops.h |
| 226 | |
| 227 | // TF filter shape is [ H, W, ..., inC, outC ] |
| 228 | |
| 229 | // The input gradients are computed by a convolution of the output gradients |
| 230 | // and the filter, with some appropriate padding. See the comment at the top |
| 231 | // of conv_grad_ops.h for details. |
| 232 | PrimitiveType dtype = out_backprop_xla_shape.element_type(); |
| 233 | auto out_backprop_div = AvgPoolDivideByCount( |
| 234 | out_backprop, gradients_size, kernel_size, stride, spatial_padding, |
| 235 | dtype, data_format, counts_include_padding); |
| 236 | |
| 237 | // Pad the gradients in the spatial dimensions. We use the same padding |
| 238 | // as Conv2DBackpropInput. |
| 239 | PaddingConfig padding_config = MakeNoPaddingConfig(num_dims); |
| 240 | std::vector<int64> padded_gradients_size(gradients_size.begin(), |
| 241 | gradients_size.end()); |
| 242 | // First, pad the output gradients the same way as the input. The additional |
| 243 | // padding will be removed as a last step before returning the input |
| 244 | // gradients. |
| 245 | const int num_spatial_dims = num_dims - 2; |
| 246 | for (int i = 0; i < num_spatial_dims; ++i) { |
| 247 | int dim = data_format.spatial_dimension(i); |
| 248 | padded_gradients_size[dim] += |
| 249 | (spatial_padding[i].first + spatial_padding[i].second); |