import torch


def focal_loss(yhat, ytrue, alpha=0.75, gamma=2):
    """
    Computes α-balanced focal loss from FAIR
    https://arxiv.org/pdf/1708.02002v2.pdf
    Args:
        yhat (Tensor): predicted masks
        ytrue (Tensor): targets masks
        alpha (Float): weight to balance Cross entropy value
        gamma (Float): focal parameter
    output:
        loss value with `mean` reduction
    """

    # compute the actual focal loss
    focal = -alpha * torch.pow(1. - yhat, gamma) * torch.log(yhat)
    f_loss = torch.sum(ytrue * focal, dim=1)

    return torch.mean(f_loss)