Cannot use sparse COO tensor as adjacency matrix

I am trying to use a sparse tensor as adjacency matrix for graph algorithms. Here is the code:

def makeSparseAdjMatrix(inputs, outputs, outdegree):
        create a sparse adjacency matrix for a bipartite graph, with `input` input nodes, `outdegree` out edges from overy input node
    edgelist = makeBipartiteEdgelist(inputs, outputs, outdegree)
    row, col = zip(*edgelist)
    return torch.sparse_coo_tensor(torch.tensor([row, col], dtype=torch.int32),
                                   torch.ones(len(row)), torch.Size([inputs, outputs])).coalesce()#.to_sparse_csr()

class basicNet(nn.Module):
    def __init__(self, adjacency, inputs, outputs, testNum):
        super(basicNet, self).__init__()
        self.adjacency = adjacency
        self.inputs = inputs
        self.outputs = outputs
        self.x = nn.Parameter(torch.randn(inputs))      
    def forward(self):
        return, self.adjacency, reduce='amax')

def runBasicNet(inputs, outputs, outdegree, testNum):
        run a basic net on a bipartite graph with `inputs` input nodes, `outputs` output nodes, and `outdegree` out edges from every input node
    adjacency = makeSparseAdjMatrix(inputs, outputs, outdegree)
    net = basicNet(adjacency, inputs, outputs, testNum)
    if torch.cuda.is_available():
        adjacency ='cuda:0')
        net ='cuda:0')
        print('data uses cuda')
    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
    labels = torch.tensor([random.randint(0, 5) for i in range(outputs - testNum)])
    criterion = nn.CrossEntropyLoss()
    train_mask = torch.tensor([True]*(outputs-testNum) + [False]*testNum)
    for epoch in range(10):
        result = net()
        loss = criterion(result[train_mask], labels)
        print(f"Epoch: {epoch}, Loss: {loss}")

That is the error:

NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCUDA' backend.

When trying to run the same on CPU:

NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCPU' backend.

Where is it implemented?

Hi Alex!

Pytorch’s support for sparse tensors is a work in progress. It appears
that the combination of things you need for your use case is not (yet)

From the documentation for the current stable release (2.3):

Supports both CSR and COO storage formats.


This function doesn’t support computing derivaties with respect to CSR matrices.


reduce is implemented only for CSR storage format on CPU device.

So if you need to backpropagate through adjacency (which it appears
you do), you have to use a COO matrix, but then you lose support for

Here is a related github issue, but it appears to be silent about
backpropagation through the sparse tensor.

Good luck!

K. Frank