import torch


def covariance(a: torch.Tensor, b: torch.Tensor, covariance_weights: torch.Tensor) -> torch.Tensor:
    """
    Calculates the covariance via :math:`w_{c} \\cdot a \\cdot b^T`.

    Args:
        a: The A matrix.
        b: The B matrix.
        covariance_weights: The weights to use for aggregating the covariance.
    """

    cov = a.unsqueeze(-1) * b.unsqueeze(-2)

    return (covariance_weights.view(-1, 1, 1) * cov).sum(-3)