Create a knn adjacent matrix in Pytorch

user_10126246 0
Viewed 160

In Pytorch, say I have a top-k indexing matrix P(B,N,k), a weight matrix W(B,N,N) and a target matrix A(B,N,N), I want to get a adjacent matrix that operates as the following loops:

for i in range(B):
     for ii in range(N):
         for j in range(k):
             if weighted:
                 A[i][ii][P[i][ii][j]] = W[i][ii][P[i][ii][j]]
             else:
                 A[i][ii][P[i][ii][j]] = 1

How to implement it more efficiently and concisely in Pytorch?

tensor pytorch indexing matrix
1 Answers

I think you are looking for torch.scatter_:

A.scatter_(dim=2, index=P, src=W)  # for the weighted version
A.scatter_(dim=2, index=P, src=torch.ones_like(W))  # for the un-weighted version
user_17144100
Related