would this correspond in pure pytorch? it computes inner product between the central pixel to every pixels in the surrounding:
import torch
from einops import rearrange
unfold_op = torch.nn.Unfold(k1, dilation=1, padding=k1//2, stride=1)

def pixel_cost_volume(im1, im2, unfold_op):
    """im1/2 b c h w

    -> produces b h w k**2 cost volume, each channel represents
    inner product with a different discrete shift (i,j)
    """
    b,c,h,w = im1.shape
    central_pixel = rearrange(im1, 'b c h w -> b h w c')
    neighbors = unfold_op(im2).reshape(b,-1,h,w)
    neighbors = rearrange(neighbors, 'b (c k2) h w -> b h w k2 c', c=c)
    y = torch.einsum('...i c, ...c -> ...i', neighbors, central_pixel)
    return y # (b, h, w, search_range x search_range)