| 131 | } // namespace |
| 132 | |
| 133 | XlaOp MaxPool(XlaOp operand, absl::Span<const int64> kernel_size, |
| 134 | absl::Span<const int64> stride, Padding padding, |
| 135 | const TensorFormat& data_format) { |
| 136 | XlaBuilder* b = operand.builder(); |
| 137 | return b->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 138 | TF_ASSIGN_OR_RETURN(Shape operand_shape, b->GetShape(operand)); |
| 139 | PrimitiveType dtype = operand_shape.element_type(); |
| 140 | auto max_computation = CreateScalarMaxComputation(dtype, b); |
| 141 | auto init_value = MinValue(b, dtype); |
| 142 | return ReduceWindow(operand, init_value, max_computation, kernel_size, |
| 143 | stride, padding); |
| 144 | }); |
| 145 | } |
| 146 | |
| 147 | XlaOp AvgPool(XlaOp operand, absl::Span<const int64> kernel_size, |
| 148 | absl::Span<const int64> stride, |