I have two 3D tensors X and Q of shape (5, 16, 128) on which I do cosine similarity on 2nd dim to get a (5, 16) cosine-similarity vector. I then sort this cosine-similarity vector, to get indices of most-to-least similar vectors in the original vector Q. In code
X = torch.ones((5, 16, 128)) Q = torch.randn(5, 16, 128) # dim=2 is features neg_samples = 11 pos_samples = 5 cos_sim = torch.nn.CosineSimilarity(dim=2, eps=1e-6) sim_scores = cos_sim(X, Q) # sim_scores output shape: [5,16] _, sorted_idxs = torch.topk(sim_scores, k=sim_scores.shape, sorted=True, dim=1) #sort sim_scores to get sorted indexes same shape [5, 16] #since sorted_idxs is sorted descending on similarity, we can slice to divide pos/neg samples. pos_idxs = sorted_idxs[:, :pos_samples] #gives [5, 5] most positive idxs from sorted_idxs neg_idxs = sorted_idxs[:, pos_samples:(pos_samples+neg_samples)] #gives [5, 11] most negative idxs from sorted_idxs
How do i retrieve the corresponding tensors from
Q which is of shape
[5, 16, 128] using these
Q[ pos_idxs, :] but i get output tensor of shape
[5, 5, 16, 128]. I want to slice
[5, 5, 128] and
[5, 11, 128] two separate tensors based on 2D
neg_idxs index tensors.
How do i do this in pytorch using tensor operations?