# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/vision.models.pooling.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['BlurPool', 'MaxBlurPool']

# Cell
#nbdev_comment from __future__ import annotations

from kornia.filters import BlurPool2D, MaxBlurPool2D

from ...imports import *

# Cell
def BlurPool(
    stride:int=2, # The stride size for pooling
    ks:int=3, # The kernel size for pooling
    padding:int=0, # Unused, for fastai compatibility
    ndim:int=2, # Unused, for fastai compatibility
    ceil_mode:bool=False # Unused, for fastai compatibility
) -> BlurPool2D:
    "Compute blur (anti-aliasing) and downsample a given feature map."
    assert ndim==2, "BlurPool is only implemented for 2D"
    return BlurPool2D(kernel_size=ks, stride=stride)

# Cell
def MaxBlurPool(
    stride:int=2, # The stride size for blur pooling
    ks:int=3, # The kernel size for blur pooling
    padding:int=0, # Unused, for fastai compatibility
    ndim:int=2, # Unused, for fastai compatibility
    ceil_mode:int=True, # If True, output size matches conv2d with same kernel size
    max_ks:int=2 # The kernel size for max pooling
) -> MaxBlurPool2D:
    "Compute pools and blurs and downsample a given feature map. Equivalent to `nn.Sequential(nn.MaxPool2d(...), BlurPool2D(...))`"
    assert ndim==2, "MaxBlurPool is only implemented for 2D"
    return MaxBlurPool2D(kernel_size=ks, stride=stride, ceil_mode=ceil_mode, max_pool_size=max_ks)