| 145 | } |
| 146 | |
| 147 | XlaOp AvgPool(XlaOp operand, absl::Span<const int64> kernel_size, |
| 148 | absl::Span<const int64> stride, |
| 149 | absl::Span<const std::pair<int64, int64>> padding, |
| 150 | const TensorFormat& data_format, |
| 151 | const bool counts_include_padding) { |
| 152 | XlaBuilder* b = operand.builder(); |
| 153 | return b->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 154 | TF_ASSIGN_OR_RETURN(Shape operand_shape, b->GetShape(operand)); |
| 155 | PrimitiveType dtype = operand_shape.element_type(); |
| 156 | auto init_value = Zero(b, dtype); |
| 157 | std::vector<int64> input_size(operand_shape.dimensions().begin(), |
| 158 | operand_shape.dimensions().end()); |
| 159 | const int num_dims = kernel_size.size(); |
| 160 | const int num_spatial_dims = num_dims - 2; |
| 161 | auto padding_config = MakeSpatialPaddingConfig(padding, num_spatial_dims, |
| 162 | stride, data_format); |
| 163 | auto padded_operand = Pad(operand, Zero(b, dtype), padding_config); |
| 164 | auto pooled = ComputeSums(padded_operand, init_value, kernel_size, stride, |
| 165 | data_format); |
| 166 | return AvgPoolDivideByCount(pooled, input_size, kernel_size, stride, |
| 167 | padding, dtype, data_format, |
| 168 | counts_include_padding); |
| 169 | }); |
| 170 | } |
| 171 | |
| 172 | std::vector<std::pair<int64, int64>> MakeSpatialPadding( |
| 173 | absl::Span<const int64> input_size, absl::Span<const int64> kernel_size, |