def expectation_2d(values, probabilities):
    """Calculate the expected value over values in a 2D layout.

    Args:
        values (torch.Tensor): Values for each position.
        probabilities (torch.Tensor): Probabilities for each position.

    Returns:
        The expected values.
    """

    prod = values * probabilities
    *first_dims, height, width = prod.size()
    mean = prod.view(*first_dims, height * width).sum(-1, keepdim=False)
    return mean