How to use GNN encoders with KGE models?

[Rumengol]

Hello,

I am trying to use PyG to build an KGE autoencoder with a GNN (say, SAGEConv) as encoder and a KGE model (say, DistMult) as a decoder. However, I only found shallow embedding models in the examples. Looking at the documentation, it seems heterogeneous graph samplers only support one edge type and are thus not applicable for this task (unless I have one dataloader for each edge type).

I tried to implement my own but due to the structure of HeteroData there is always a point where I process triplets sequentially, which creates a very bad bottleneck in computing time.

Am I missing something here? I find it especially strange not being able to sample different edge types in a heterogeneous graph. If not, any pointer on how to reconstruct edge_index_dict from a batch would be appreciated.

Thank you for your help!