Consider X[0] : [0, 1, 2, 0, 1, 0]. What I want to build an adjacency matrix such that:

A[0, 1] = 1
A[1, 2] = 1
A[2, 0] = 1
A[0, 1] = 1
A[1, 0] = 1

The code I shared works, however it’s very slow since I’m looping over each element in the batch. I’d like to vectorize it as much as possible. Thanks!