def METHOD_NAME(
    labels: torch.Tensor, predictions: torch.Tensor, weights: torch.Tensor
) -> Dict[str, torch.Tensor]:
    return {
        "weighted_sum": (predictions * weights).sum(dim=-1),
        "weighted_num_samples": weights.sum(dim=-1),
    }