import torch


def _mean_top_sim(sim: torch.Tensor, k: int, dim: int) -> torch.Tensor:
    """Compute the mean similarity of the top-k matches along an axis."""
    return sim.topk(k=k, dim=dim, largest=True, sorted=False)[0].mean(dim=dim).unsqueeze(dim=dim)