Laplacian positional encoding in a pytorch batch
I am trying to recode the laplacian positional encoding for a graph model in pytorch. A valid encoding in numpy can be found at https://docs.dgl.ai/en/0.9.x/_modules/dgl/transforms/functional.html#laplacian_pe .
I think I have managed to make an equivalent encoding to numpy in pytorch, but for performance issues I would like that function to be able to work with batches of data.
That is, the following function works with the parameters
with the form adj[N, N], degrees[N, N] and topk as an integer, where N is the number of nodes in the network.
def _laplacian_positional_encoding_th(self, adj, degrees, topk):
    number_of_nodes = adj.shape[-1].
    #degrees = th.clip(degrees, 0, 1) # not multigraph
    assert topk &lt; number_of_nodes
    
    # Laplacian
    D = th.diag(degrees**-0.5)
    B = D * adj * D
    L = th.eye(number_of_nodes).to(B.device) * B

    # Eigenvectors
    EigVal, EigVec = th.linalg.eig(L)
    idx = th.argsort(th.real(EigVal)) # increasing order
    EigVal, EigVec = th.real(EigVal[idx]), th.real(EigVec[:,idx])

    # Only select [1,topk+1] EigenVectors as L is symmetric (Spectral decomposition)
    out = EigVec[:,1:topk+1]
    return out

However, when I try to perform the same efficient operations in batch form, I cannot code it. That is, the idea is that the parameters can come in the form adj[B, N, N], degrees[B, N, N] and topk as integer, B being the number of data in the batch.

How about:
def _laplacian_positional_encoding_th(self, adj, degrees, topk):
        number_of_nodes = adj.shape[-1]
        assert topk &lt; number_of_nodes

        D = th.clip(degrees, 0, 1) # not multigraph
        B = D @ adj @ D
        L = th.eye(number_of_nodes).to(B.device)[None, ...] - B 

        # Eigenvectors
        EigVal, EigVec = th.linalg.eig(L)
        idx = th.argsort(th.real(EigVal)) # increasing order

        out = th.real(th.gather(EigVec, dim=-1, index=idx[..., None]))
        return out

See th.diag_embed for creating a batch of diagonal matrices, and th.gather for selecting the right columns of EigVec according to the sorted indices.
Updata:
If you want to extract the topk vectors:
_, topk = th.topk(EigVal.real, k=5)  # get the top 5
out = th.gather(EigVec.real, dim=-1, index=topk[:, None, :].expand(-1, EigVec.shape[1], -1))

