Form triplets from a batch of Pytorch tensor embeddings

form_triplets.py

def form_triplets(inA, inB):
        
    '''
    Form triplets from two tensors of embeddings. It is assumed that the embeddings at corresponding batch positions are similar
    and all other batch positions are dissimilar 
    
    i.e inA[i] ~ inB[i] and inA[i] !~ inB[j] for all i =! j
    '''
    
    b, emb_size = inA.shape
    perms = b**2
    
    labels = [0]*perms; sim_idxs = [(0 + i*b) + i for i in range(b)]
    for idx in sim_idxs:
        labels[idx] = 1
    labels = torch.Tensor(labels)
        
    labels = labels.type(torch.BoolTensor).to(inA.device)
    anchors = inA.repeat(b, 1)[~labels]
    negatives = torch.cat([inB[i,:].repeat(b,1) for i in range(b)])[~labels]
    positives = inB.repeat(b, 1)[~labels]

    return(anchors, positives, negatives)