title: Create a knn adjacent matrix in Pytorch
body: <p>In Pytorch, say I have a top-k indexing matrix <code>P(B,N,k)</code>, a weight matrix <code>W(B,N,N)</code> and a target matrix <code>A(B,N,N)</code>, I want to get a adjacent matrix that operates as the following loops:</p>
  <pre><code>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
  </code></pre>
  <p>How to implement it more efficiently and concisely in <strong>Pytorch</strong>?</p>