| 170 | } |
| 171 | |
| 172 | std::vector<std::pair<int64, int64>> MakeSpatialPadding( |
| 173 | absl::Span<const int64> input_size, absl::Span<const int64> kernel_size, |
| 174 | absl::Span<const int64> stride, Padding padding, |
| 175 | const TensorFormat& data_format) { |
| 176 | const int num_spatial_dims = kernel_size.size() - 2; |
| 177 | std::vector<int64> input_spatial_dimensions; |
| 178 | std::vector<int64> kernel_size_spatial_dimensions; |
| 179 | std::vector<int64> stride_spatial_dimensions; |
| 180 | CHECK_EQ(data_format.num_spatial_dims(), num_spatial_dims) |
| 181 | << "Invalid number of spatial dimensions in data format specification"; |
| 182 | for (int i = 0; i < num_spatial_dims; ++i) { |
| 183 | int dim = data_format.spatial_dimension(i); |
| 184 | input_spatial_dimensions.push_back(input_size[dim]); |
| 185 | kernel_size_spatial_dimensions.push_back(kernel_size[dim]); |
| 186 | stride_spatial_dimensions.push_back(stride[dim]); |
| 187 | } |
| 188 | return MakePadding(input_spatial_dimensions, kernel_size_spatial_dimensions, |
| 189 | stride_spatial_dimensions, padding); |
| 190 | } |
| 191 | |
| 192 | XlaOp AvgPoolGrad(XlaOp out_backprop, absl::Span<const int64> gradients_size, |
| 193 | absl::Span<const int64> kernel_size, |