How can I see source code or explanation of "torch_sparse import SparseTensor"?

Question

I am studying some source codes from PytorchGeometric.

Actually I am really finding from torch_sparse import SparseTensor in Google, to get how to use SparseTensor.

But there is nothing I can see explanation. I saw many documents about COO,CSR something like that, but how can I use SparseTensor?

I read : https://pytorch.org/docs/stable/sparse.html# but there is nothing like SparseTensor.

Thank you in advance :)

Answer 1

I just had the same problem and stumbled upon your question, so I will just detail what I did here, maybe it helps someone. I think the main confusion results from the naming of the package. SparseTensoris from torch_sparse, but you posted the documentation of torch.sparse. The first is an individual project in the pytorch ecosystem and a part of the foundation of PyTorch Geometric, but the latter is a submodule of the actual official PyTorch package.

So, looking at the right package (torch_sparse), there is not much information about how to use the SparseTensor class there (Link). If we go to the source code on the other hand (Link) you can see that the class has a bunch of classmethods that you can use to genereate your own SparseTensor from well documented pytorch classes.

In my case, all I needed was a way to feed the RGCNConvLayer with just one Tensor including both the edges and edge types, so I put them together with the following line:

edge_index = SparseTensor.from_edge_index(edge_index, edge_types)

If you, however, already have a COO or CSR Tensor, you can use the appropriate classmethods instead.