Create a knn adjacent matrix in Pytorch

Question

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 adjace