pyg-team / smore Goto Github PK

The overall goal is to fit SMORE into the PyG interface, and to re-use as much of PyG's functionality as possible. This involves:

• Providing datasets as part of a InMemoryDataset
• Integrate SparseEmbedding functionality with asychronous access directly in PyG (can be later used for integration with GNNAutoScale as well)
• Clean-up models to make use of MessagePassing whenever possible
• Out-source evaluation metrics to torch_geometric.utils or torch_geometric.metrics
• Potentially integrate with GraphGym once multi-GPU support for GraphGym lands

In general, I'm in favor of storing KGs as pure PyTorch tensors, and utilize the torch.extension package for C++/CUDA integration (similar to how torch-scatter and torch-sparse works). This should also allow us to have @torch.jit.script support for C++ extensions.

We can store KGs in CSR representation based on head type. Relations can be sorted based on relation id inside each head id and can be then effectively queried via binary search.