"""Utils for criterion."""
import torch
import torch.nn.functional as F


def normalize(x, axis=-1):
    """Performs L2-Norm."""
    num = x
    denom = torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12
    return num / denom


# Source : https://github.com/earhian/Humpback-Whale-Identification-1st-/blob/master/models/triplet_loss.py
def euclidean_dist(x, y):
    """Computes Euclidean distance."""
    m, n = x.size(0), y.size(0)
    xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
    yy = torch.pow(x, 2).sum(1, keepdim=True).expand(m, m).t()
    dist = xx + yy - 2 * torch.matmul(x, y.t())

    dist = dist.clamp(min=1e-12).sqrt()

    return dist


def cosine_dist(x, y):
    """Computes Cosine Distance."""
    x = F.normalize(x, dim=1)
    y = F.normalize(y, dim=1)
    dist = 2 - 2 * torch.mm(x, y.t())
    return dist

[((28, 8, 28, 29), 'torch.nn.functional.normalize', 'F.normalize', (), '', True, 'import torch.nn.functional as F\n'), ((29, 8, 29, 29), 'torch.nn.functional.normalize', 'F.normalize', (), '', True, 'import torch.nn.functional as F\n'), ((9, 12, 9, 48), 'torch.norm', 'torch.norm', (), '', False, 'import torch\n'), ((17, 9, 17, 24), 'torch.pow', 'torch.pow', ({(17, 19, 17, 20): 'x', (17, 22, 17, 23): '2'}, {}), '(x, 2)', False, 'import torch\n'), ((18, 9, 18, 24), 'torch.pow', 'torch.pow', ({(18, 19, 18, 20): 'x', (18, 22, 18, 23): '2'}, {}), '(x, 2)', False, 'import torch\n')]