How to implement X^T A X effectively in torch_geometric?

[jiaxilv]

I would be grateful if you could show me how to implement formula like X^T A X effectively.
Thanks:)

[wsad1]

Assuming A is an adjacency matrix a good way to do this would be

index,value=torch_geometric.utils.sparse.dense_to_sparse(A) 
ax=torch_sparse.spmm(index,value,n,f,X) #n=num nodes, f=node dimension 
xax=torch.matmul(X.T,ax)

check out torch_sparse.spmm for more details.

[jiaxilv]

hey！Thank you for your reply！
Suppose data is torch_geometric.data.Batch like Batch(batch=[866], edge_index=[2, 4284], x=[866, 89], y=[3]).
If I use the above code directly, I get a matrix shape of (89,89),which loses the batch information.
It is easy to calculate the results using a loop serial, but it took too long.
`
class model(nn.Module):
.........

def forward(self,x,edge_index,batch):

      num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0)

      batch_size, max_num_nodes = batch.max().item()+1, num_nodes.max().item()

      out = []

      for i in range(batch_size):

          x_s = x[batch == i]

          x_i_out = #compute X^T A X for a single graph

          out.append(x_i_out)

      return torch.stack(out,dim = 0)

`
I just wonder how can I compute it more efficiently when data is torch_geometric.data.Batch.
Looking for your reply:)

[rusty1s]

This is really tricky to implement in a sparse-scenario, since we do not support batch-wise multiplication with sparse tensors. Is it applicable to you to convert your adjacency matrix to a dense one first?

[jiaxilv]

I know what u mean.you mean

def forward(x,edge_index,batch):
x = to_dense_batch(x,batch)[0]
adj = to_dense_adj(edge_index,batch)
out = torch.matmul(adj,x)
out = torch.matmul(x.transpose(1,2),out)
return out

This is a good way, while more memory is needed.
But maybe everything is ok on 2080Ti .
Thank you for you reply and your torch_geometric.
May u have a good day:)

[flandolfi]

Hi! Notice that, assuming A is a SparseTensor and X a Tensor, until

Y = A @ X 

the results are still correct (batchwise). The tricky part arrives when you try to do X.T @ Y, which is indeed not correct.
You can do as follows:

Y = A @ X
X, mask = to_dense_batch(X, batch)
Y, _ = to_dense_batch(Y, batch)  # the mask is the same
X = torch.bmm(X.transpose(-1, -2), Y)

The difference is that you do not compute a dense adjacency matrix, so you use up to O(bnh) and not O(bn²) space, where b is the batch size, n the maximum graph size in the batch and h the data dimension (X.size(1)).