Cannot use sparse COO tensor as adjacency matrix

alexj: <p>Hi,<br>
I am trying to use a sparse tensor as adjacency matrix for graph algorithms. Here is the code:</p>
<pre><code class="lang-auto">def makeSparseAdjMatrix(inputs, outputs, outdegree):
    """_summary_
        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 torch.sparse.mm(self.x, self.adjacency, reduce='amax')

def runBasicNet(inputs, outputs, outdegree, testNum):
    """_summary_
        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 = adjacency.to('cuda:0')
        net = net.to('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):
        optimizer.zero_grad()
        result = net()
        loss = criterion(result[train_mask], labels)
        loss.backward()
        optimizer.step()
        print(f"Epoch: {epoch}, Loss: {loss}")
</code></pre>
<p>That is the error:</p>
<pre><code class="lang-auto">NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCUDA' backend.
</code></pre>
<p>When trying to run the same on CPU:</p>
<pre><code class="lang-auto">NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCPU' backend.
</code></pre>
<p>Where is it implemented?</p>

KFrank: <p>Hi Alex!</p>
<aside class="quote no-group quote-modified" data-username="alexj" data-post="1" data-topic="204367" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alexj/48/70445_2.png" class="avatar"> alexj:</div>
<blockquote>
<p>I am trying to use a sparse tensor as adjacency matrix<br>
…</p>
<pre><code class="lang-auto">        return torch.sparse.mm(self.x, self.adjacency, reduce='amax')
</code></pre>
<pre><code class="lang-auto">NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCUDA' backend.
</code></pre>
<p>Where is it implemented?</p>
</blockquote>
</aside>
<p>Pytorch’s support for sparse tensors is a work in progress.  It appears<br>
that the combination of things you need for your use case is not (yet)<br>
implemented.</p>
<p>From the <a href="https://pytorch.org/docs/2.3/generated/torch.sparse.mm.html" rel="noopener nofollow ugc">sparse.mm documentation</a> for the current stable release (2.3):</p>
<blockquote>
<p>Supports both CSR and COO storage formats.</p>
</blockquote>
<p>but</p>
<blockquote>
<p>This function doesn’t support computing derivaties with respect to CSR matrices.</p>
</blockquote>
<p>and</p>
<blockquote>
<p><code>reduce</code> is implemented only for CSR storage format on CPU device.</p>
</blockquote>
<p>So if you need to backpropagate through <code>adjacency</code> (which it appears<br>
you do), you have to use a COO matrix, but then you lose support for<br>
<code>reduce</code>.</p>
<p>Here is <a href="https://github.com/pytorch/pytorch/issues/105790" rel="noopener nofollow ugc">a related github issue,</a> but it appears to be silent about<br>
backpropagation through the sparse tensor.</p>
<p>Good luck!</p>
<p>K. Frank</p>