Pytorch implementation of sigsoftmax - https://arxiv.org/pdf/1805.10829.pdf

sigsoftmax.py

def logsigsoftmax(logits):
    """
    Computes sigsoftmax from the paper - https://arxiv.org/pdf/1805.10829.pdf
    """
    max_values = torch.max(logits, 1, keepdim = True)[0]
    exp_logits_sigmoided = torch.exp(logits - max_values) * torch.sigmoid(logits)
    sum_exp_logits_sigmoided = exp_logits_sigmoided.sum(1, keepdim = True)
    log_probs = logits - max_values + torch.log(torch.sigmoid(logits)) - torch.log(sum_exp_logits_sigmoided)
    return log_probs

@yottabytt
See: https://en.wikipedia.org/wiki/LogSumExp

I think you can also get the log_probs by:

sigmoid_logits = logits.sigmoid().log()
sigsoftmax_logits = logits + sigmoid_logits
return sigsoftmax_logits.log_softmax()