def _pool(x, kernel_size, stride, dilation, reduction, return_indices=False):
    """Implement pooling by "manually" extracting patches using `unfold`

    Parameters
    ----------
    x : (batch, *spatial)
    kernel_size : (dim,) int
    stride : (dim,) int
    dilation : (dim,) int
    reduction : callable
        This function should collapse the last dimension of a tensor
        (..., K) tensor -> (...)
    return_indices : bool

    Returns
    -------
    x : (batch, *spatial_out)

    """
    kernel_size = [(sz-1)*dl + 1 for sz, dl in zip(kernel_size, dilation)]
    for d, (sz, st, dl) in enumerate(zip(kernel_size, stride, dilation)):
        x = x.unfold(dimension=d + 1, size=sz, step=st)
        if dl != 1:
            x = x[..., ::dl]
    dim = len(kernel_size)
    patch_shape = x.shape[dim+1:]
    x = x.reshape((*x.shape[:dim+1], -1))
    if return_indices:
        x, ind = reduction(x, return_indices=True)
        return x, ind
    else:
        x = reduction(x)
        return x