def marginal_log_loss(logits, is_correct):
                """Loss based on the negative marginal log-likelihood."""

                def mask_to_score(mask, dtype=torch.float32):
                    return (1.0 - mask.type(dtype)) * torch.finfo(dtype).min

                # []
                log_numerator = torch.logsumexp(logits + mask_to_score(is_correct, dtype=logits.dtype), dim=-1)
                log_denominator = torch.logsumexp(logits, dim=-1)
                return log_denominator - log_numerator
Loss based on the negative marginal log-likelihood.
dpkd/transformers/src/transformers/models/realm/modeling_realm.py