I have a matrix of embeddings of size `N*d`

(N inputs each having an embedding of size d) and I have a graph adjacency matrix of size `N*N`

. Assuming the adjacency matrix is as follows (N=4):

[[1,1, 0, 0],

[0, 0, 1, 0],

[1, 0, 0, 0],

[1, 0, 1, 0]]

And the embedding matrix is as follows (d=3):

[[1, 2, 3],

[4, 5, 6],

[7, 8, 9],

[10, 11, 12]]

I want to get the following output:

[[18, 21, 24],

[1, 2, 3],

[14, 16, 18],

[0, 0, 0, 0]]

The output is computed as follows: since the first column of the adjacency matrix is 1 at rows 0, 2, and 3, we sum the embeddings at rows 0, 2, 3 and that becomes the first row of the output ([18, 21, 24]). Similarly, since the second column of the adjacency matrix is 1 only at row 0, the second row of the output is just equivalent to row 0 of the embedding. Since the third column of the adjacency matrix is 1 at rows 1 and 3, the third row of the output is the sum of the rows 1 and 3 of the embedding. And since the last column of the adjacency matrix is all zeros, the last row of the output is all zeros.

Canâ€™t figure out how to do it without using a for loop. Thanks for any helps in advance.