import numpy as np
from sklearn.feature_extraction.image import extract_patches

def conv2d(inputs, filters):
    """
    Args:
        inputs (np.ndarray): NHWC
        filters (np.ndarray):
            with shape [filter_height, filter_width, in_channels, out_channels]
    """

    kH, kW, inC, outC = filters.shape
    patches = extract_patches(images, (1, kH, kW, 1))
    patches = patches.reshape((*patches.shape[:3], -1))

    kernel = filters.reshape((kH*kW*inC, outC))

    return patches @ filters.reshape((kH*kW*inC, outC))

if __name__ == "__main__":

    B = 4
    H, W = 11, 11
    inC = 3
    outC = 6
    kH, kW = 3,3

    images = np.random.randint(0, 2, (B,H,W, inC))
    kernel = np.random.random((kH,kW,inC,outC))

    h = conv2d(images, kernel)