submission 512856

import torch
import torch.nn.functional as F
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    input_tensor, kernel, output = data
    
    batch, channels, height, width = input_tensor.shape
    out_channels, in_channels, kh, kw = kernel.shape
    
    out_h = height - kh + 1
    out_w = width - kw + 1
    
    # im2col: (batch, in_channels*kh*kw, out_h*out_w)
    input_unfolded = F.unfold(input_tensor, kernel_size=(kh, kw), stride=1, padding=0)
    
    # kernel: (out_channels, in_channels*kh*kw)
    kernel_matrix = kernel.view(out_channels, -1)
    
    # matmul broadcast: (out_channels, K) x (batch, K, out_h*out_w) -> (batch, out_channels, out_h*out_w)
    result = torch.matmul(kernel_matrix, input_unfolded)
    
    output.copy_(result.view(batch, out_channels, out_h, out_w))
    return output