[Roadmap] PyTorch SparseTensor Support 🚀

[EdisonLeeeee]

🚀 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 (Add spmm and is_torch_sparse_tensor #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 (Add spmm and is_torch_sparse_tensor #5906, Added mean reduction to utils.spmm for PyTorch Sparse Tensor #6868, PyTorch Sparse Tensor support: HANConv, GATv2Conv, HGTConv, GMMConv, GPSConv, and RGATConv #6932).
• is_sparse_tensor: check for either torch.sparse.Tensor or torch_sparse.SparseTensor (Add is_sparse and to_torch_coo_tensor #6003).
• to_torch_coo_tensor: convert edge_index to torch.sparse.Tensor (coo format) (Add is_sparse and to_torch_coo_tensor #6003).
• Support torch.sparse.Tensor input of add_self_loop, remove_self_loop, and maybe_num_nodes (Added PyTorch Sparse Tensor support for remove_self_loops, add_self_loops, and maybe_num_nodes #6847)

torch_geometric.nn.*

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

torch_geometric.transforms.*

• ToSparseTensor (Added PyTorch Sparse Tensor support for ToSparseTensor #6930)

[EdisonLeeeee]

So far I can think of a way to support torch SparseTensor with the least effort in MessagePassing is:

def hook(model, input):
    edge_index, size, kwargs = input
    if is_torch_coo_tensor(edge_index):
        adj = edge_index
        edge_index = adj._indices()
        kwargs['edge_weight'] = adj._values()
    return edge_index, size, kwargs

register_propagate_forward_pre_hook(hook)

we can register such a hook in __init__ of MessagePassing, WDYT? @rusty1s

(Update) This is not a perfect solution as we cannot enjoy the benefits of PyTorch SparseTensor in this way.

[rusty1s]

I think this would only be a temporary solution. Ultimately, we want direct support for it similar to SparseTensor. Wondering what are the main challenges to support this?

[EdisonLeeeee]

IMO, there are some challenges to support this:

1. It only supports sum aggregation when using PyTorch SparseTensor. For other cases including advanced aggregations, we still need to convert it to the form of (edge_index, edge_weight).
2. PyTorch SparseTensor does not support spspmm or mul with broadcasting, which makes it challenging to implement some operations such as gcn_norm (e.g., $A\cdot D^{-\frac{1}{2}}$ ) efficiently.
3. As mentioned above, since PyTorch SparseTensor is also an instance of torch.Tensor, we should change the condition isinstance(edge_index, Tensor) if we want to support it. There should be some auxiliary functions like is_edge_index, is_torch_coo_tensor, or something else.

Please correct me if something is missing.

[rusty1s]

These are good thoughts. Appreciate it. Let me think.

1. The only way I can think of supporting this is to provide a spmm function in PyG that supports both torch_sparse and PyTorch. For PyTorch, we can error out in case aggr != "sum" or aggr != "mean".
2. Yes, this is a real problem. We would need to provide our own implementation for this by working on the indices and values directly. Happy to move this to follow-up PRs and work on general MessagePassing integration first
3. Yes, this is expected, but IMO not a blocker, right?

[EdisonLeeeee]

You are right. Thanks for making it clear. Will update the roadmap correspondingly and make some PRs soon :)

[EdisonLeeeee]

Just go back from vacation :) I've updated the roadmap and corresponding TODO list. I plan to support them in the following weeks.

[rusty1s]

Nice to have you back :)