import torch


def get_anchor_positive_triplet_mask(labels):
    """
    Return a 2D mask where mask[a, p] is True iff a and p are distinct and have same label.
    :param labels: tensor of shape (batch_size, )
    :return: tensor of shape (batch_size, batch_size)
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # check that i and j are distinct
    indices_not_equal = torch.eye(labels.shape[0]).to(device).byte() ^ 1

    # check if labels[i] == labels[j]
    labels_equal = torch.unsqueeze(labels, 0) == torch.unsqueeze(labels, 1)

    # combine the two masks
    mask = indices_not_equal * labels_equal

    return mask