Support for CSR Tensor with NN layers

[vgutta]

🐛 Describe the bug

When I try to pass a CSR tensor to a forward pass for a NN it outputs NaN.

Here is the NN:

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(292, 8)
        self.conv2 = GCNConv(8, 7)

    def forward(self, data):
         x, edge_index = data.x.to_dense(), data.edge_index
         x = self.conv1(x, edge_index)
         x = F.relu(x)
         x = F.dropout(x, training=self.training)
         x = self.conv2(x, edge_index)
         return F.log_softmax(x, dim=1)

Are there any specific reasons why CSR tensors are not supported yet?

Versions

PyTorch 1.10

cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @saketh-are

[amjames]

Could you provide a full reproduction for this?

I tried this out and as you call to_dense in forward there shouldn't be any issue. Without that transformation there are issues with the backward pass, but not forward.

[vgutta]

@amjames Are there any recommended solutions for dealing errors during backward pass for data in sparse format?

Example errors found here

[amjames]

@vgutta I see errors here at the backward step, these are known issues and we are working toward better autograd support for sparse layouts.

I am curious specifically about your original report that passing a sparse CSR tensor to the forward pass results in NaN. This is unexpected and I could not reproduce that result. Could you point out where that is observed? I did not see it in the linked notebook.

[cpuhrsch]

@amjames @vgutta - The issue cites PyTorch version 1.10. Does this persist on a recent nightly? Thanks.

[vgutta]

@vgutta I see errors here at the backward step, these are known issues and we are working toward better autograd support for sparse layouts.

I am curious specifically about your original report that passing a sparse CSR tensor to the forward pass results in NaN. This is unexpected and I could not reproduce that result. Could you point out where that is observed? I did not see it in the linked notebook.

@amjames Sorry forgot to mention in my previous reply that the original issue may have had more to do with the dataset and model I was using rather than a PyTorch issue.

@cpuhrsch I think as far as I can tell the backward pass does not support sparse layouts