import torch


def logsumexp_masked(risk_scores,
                     riskset):                     
    """Compute logsumexp across `axis` for entries where `mask` is true."""
    risk_score_masked = torch.mul(risk_scores, riskset).float()
    # for numerical stability, substract the maximum value
    # before taking the exponential
    amax = torch.max(risk_score_masked, 1)[0]
    risk_score_shift = risk_score_masked.sub(amax)
    exp_masked = torch.mul(risk_score_shift.exp(), riskset)
    exp_sum = torch.sum(exp_masked,1)
    output = amax + torch.log(exp_sum)      
    return output