[Roadmap] PyTorch SparseTensor Support 🚀

🚀 The feature, motivation and pitch

PyG currently accepts torch.LongTensor: edge_index and torch_sparse.SparseTensor: adj in MessagePassing, which limits its flexibility for users who use native PyTorch SparseTensor. As such, an additional step is required for them to convert one to another.

The goal of this roadmap is to track the integration of native PyTorch SparseTensor support into PyG. After integration, the edge_index in MessagePassing can also accept PyTorch SparseTensor while allowing backpropagation like torch_sparse.SparseTensor.

General Roadmap

• Add PyTorch SparseTensor support for torch_geometric.nn.MessagePassing and corresponding layers inherited from it
• Add PyTorch SparseTensor support for torch_geometric.transforms
• Add PyTorch SparseTensor support for torch_geometric.loader
• Replace torch_sparse.SparseTensor based implementations with PyTorch functionality
• ...

Implementations

Utility functions

• is_torch_sparse_tensor function. Since both strided (dense) or sparse_coo tensors are instances of torch.Tensor, there should be a function to distinguish between two different inputs, i.e., LongTensor edge_index and SparseTensor adj (#5906).
• spmm function: sparse-dense matrix multiplication supporting both torch_sparse and PyTorch SparseTensor. For PyTorch SparseTensor, currently only sum and mean aggregations are allowed (#5906, #6868, #6932).
• is_sparse_tensor: check for either torch.sparse.Tensor or torch_sparse.SparseTensor (#6003).
• to_torch_coo_tensor: convert edge_index to torch.sparse.Tensor (coo format) (#6003).
• Support torch.sparse.Tensor input of add_self_loop, remove_self_loop, and maybe_num_nodes (#6847)

torch_geometric.nn.*

• MessagePassing (#5944, #6633)
• GCNConv (#6033)
• AGNNConv (#6874)
• APPNP (#6874)
• AntiSymmetricConv (#6633)
• ARMAConv
• CGConv (#6633)
• ChebConv
• ClusterGCNConv (#6874)
• DNAConv (#6748)
• EdgeConv (#6748)
• EGConv (#6748)
• FAConv (#6748)
• FeaStConv (#6874)
• FiLMConv
• FusedGATConv
• GATConv (#6897)
• GatedGraphConv (#6897)
• GATv2Conv (#6932)
• GCN2Conv (#6897)
• GENConv (#6897)
• GeneralConv
• GINConv (#6532)
• GMMConv (#6932)
• GPSConv (#6932)
• GraphConv (#6532)
• GravNetConv
• HANConv (#6932)
• HEATConv
• HeteroConv
• HGTConv (#6932)
• HypergraphConv
• LEConv (#6936)
• LGConv (#6936)
• MFConv
• NNConv (#6932)
• PANConv (#6936)
• PDNConv
• PNAConv
• PointNetConv (#6937)
• PointGNNConv (#6937)
• PointTransformerConv (#6937)
• PPFConv (#6937)
• ResGatedGraphConv (#6937)
• RGATConv (#6932)
• RGCNConv & FastRGCNConv
• SAGEConv (#6532, #6874)
• SGConv (#6514)
• SignedConv (#6936)
• SplineConv
• SSGConv (#6514)
• SuperGATConv
• TAGConv (#6514)
• TransformerConv (#6633)
• WLConvContinuous
• WLConv (#6936)
• XConv

torch_geometric.transforms.*

• ToSparseTensor (#6930)