# Cosine similarity on 3D tensors and Filtering

Hi All,
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[1], 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 `pos_idxs` shape `[5,5]` and `neg_idxs` shape `[5,11]`?

I tried `Q[ pos_idxs, :]` but i get output tensor of shape `[5, 5, 16, 128]`. I want to slice `Q` into `[5, 5, 128]` and `[5, 11, 128]` two separate tensors based on 2D `pos_idxs` and `neg_idxs` index tensors.

How do i do this in pytorch using tensor operations?

thanks