Graph-dependent matrix multiplication in a batch

[Nokimann]

Let's assume that there are n graphs in a batch.

I'm trying to multiply pos in a batch with n 3x3 matrices corresponding to each graph where matrices are different from each other with different graphs.

How can we do this matrix multiplication without repeat_interleave function?

[Nokimann]

I found segment_matmul in source of HeteroLinear seems to work. But, does it efficiently work?

Then, how can I load segment_matmul?

And does it work like the below?

pos.unsqueeze(1) @ matrix[batch]

[rusty1s]

Do you have PyG 2.2 installed?

[Nokimann]

@rusty1s

$ conda list | grep pyg
pyg                       2.2.0           py310_torch_1.13.0_cu116    pyg

Maybe, I have to reinstall my environment?

[rusty1s]

My bad, you need to look at the source code of PyG 2.2: https://pytorch-geometric.readthedocs.io/en/stable/modules/nn.html#torch_geometric.nn.dense.HeteroLinear

[Nokimann]

So, it is not implemented yet? hope to update soon...

[rusty1s]

It is integrated, the imports just changed from PyG 2.2 to master. Instead of doing from torch_geometric.typing import pyg_lib, you would just need to write import pyg_lib