import torch


def correlate_local(sparse_descriptors: torch.Tensor, dense_descriptors: torch.Tensor):
    """Compute local correlation maps.
    Args:
        * sparse_features: Sparse descriptors of size [N x C]
        * dense_features: Batch of local descriptor maps of size [B x C x H x W]
    Returns:
        * local_cmaps: The concatenated local correspondence maps, of size [B x H x W]
    """
    batch, channels, height, width = dense_descriptors.shape
    local_cmaps = torch.bmm(
        sparse_descriptors.unsqueeze(1), dense_descriptors.view(batch, channels, -1)
    ).view(batch, height, width)
    return local_cmaps