Form all pairwise examples from two Pytorch Tensors, along with a similarity label

form_pairs.py

def form_pairs(inA, inB):
    
    '''
    Form pairs 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)
    
    return(inA.repeat(b, 1), torch.cat([inB[i,:].repeat(b,1) for i in range(b)]), labels.type(torch.LongTensor).to(inA.device))