# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import itertools

import numpy as np

from megengine import Parameter, tensor
from megengine.module import AvgPool2d, MaxPool2d


def test_avg_pool2d():
    def test_func(
        batch_size,
        in_channels,
        out_channels,
        in_height,
        in_width,
        kernel_size,
        stride,
        padding,
    ):
        pool = AvgPool2d(kernel_size, stride=stride, padding=padding, mode="average")
        inp = np.random.normal(
            size=(batch_size, in_channels, in_height, in_width)
        ).astype(np.float32)
        out_height = (in_height + padding * 2 - kernel_size) // stride + 1
        out_width = (in_width + padding * 2 - kernel_size) // stride + 1
        out = pool(
tensor(inp)