utility to convert adjacency list to edge list

[abhijeetdhakane]

I am dealing with a size of array [B, P, neigh_indx], where P is an index and neigh_indx are the neighboring indices. However, PyG required it in COO format. Its possible to convert [B, P, neigh_indx] to [B, 2, P*neigh_indx] using for loops (not a right choice). I think optimized utility based on torch tensor will be useful.

Example:
"Shape : [B, P, neigh_indx]"
[[[3,5,6,7],
[5,6,8,9]]]

PyG required format
"Shape : [B, 2, P*neigh_indx]
[[[0, 0, 0, 0, 1, 1, 1, 1],
[3, 5, 6, 7, 5, 6, 8, 9]]]

[rusty1s]

Your indices will already give you the destination node indices after flattening. The source node indices can be generated via torch.arange(0, B*P).view(-1, 1).repeat(1, neigh_idx).view(-1).