import torch


def calc_output_to_cum_min_cov(
    cross_covariance: torch.Tensor,
    prob_is_min: torch.Tensor,
) -> torch.Tensor:
    """
    Given two sets of jointly Gaussian random variables [Y_1, ..., Y_d] and [Z_1, ..., Z_m],
    calculate the approximate covariance between min(Y_1, ..., Y_d) and [Z_1, ..., Z_m].

    Args:
        cross_covariance: Covariance between Gaussian variables [Y_1, ..., Y_d] and [Z_1, ..., Z_m]
            of shape [batch_dim..., d, m]
        prob_is_min: (approximate) probability that each of [Y_1, ..., Y_d] is the
            minimum. Shape [batch_dim..., d]

    Returns:
        torch.Tensor: Approximate covariance between min(Y_1, ..., Y_d) and [Z_1, ..., Z_m].
            Shape [batch_dim..., m]
    """
    return (cross_covariance * prob_is_min[..., None]).sum(dim=-2)