| 106 | } |
| 107 | |
| 108 | XlaOp AvgPoolDivideByCount(XlaOp pooled, absl::Span<const int64> input_size, |
| 109 | absl::Span<const int64> window_dimensions, |
| 110 | absl::Span<const int64> window_strides, |
| 111 | absl::Span<const std::pair<int64, int64>> padding, |
| 112 | PrimitiveType dtype, const TensorFormat& data_format, |
| 113 | bool counts_include_padding) { |
| 114 | if (counts_include_padding) { |
| 115 | // If counts include padding, all windows have the same number of elements |
| 116 | // contributing to each average. Divide by the window size everywhere to get |
| 117 | // the average. |
| 118 | int64 window_size = |
| 119 | std::accumulate(window_dimensions.begin(), window_dimensions.end(), 1, |
| 120 | [](int64 a, int64 b) { return a * b; }); |
| 121 | auto divisor = ConstantR0WithType(pooled.builder(), dtype, window_size); |
| 122 | |
| 123 | return pooled / divisor; |
| 124 | } else { |
| 125 | return AvgPoolDivideByCountWithGeneralPadding(pooled, dtype, input_size, |
| 126 | padding, window_dimensions, |
| 127 | window_strides, data_format); |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | } // namespace |
| 132 |
no test coverage detected