import torch


def bpr_loss(pos, neg, b=0.0, collapse=True):
    """
    Usual BPR loss that penalizes when negative score is greater than positive one.

    Parameters
    ----------
    pos: torch.tensor
        positive score

    neg: torch.tensor
        negative scores. Can be a batch of negatives.

    b: float
        desired boundary between positive and negative examples

    collapse: bool
        If True collapse batch with mean.

    Returns
    -------
    torch.tensor
    """
    res = torch.sigmoid(neg - pos + b)
    if collapse:
        res = res.mean()
    return res