Issue with gradient of sparse matrix using cuda

dwg: <p>Hello everyone,</p>
<p>I can replicate the code on <a href="https://pytorch.org/docs/stable/generated/torch.sparse.mm.html#torch.sparse.mm" class="inline-onebox" rel="noopener nofollow ugc">torch.sparse.mm — PyTorch 1.9.0 documentation</a> using CPU. However, when I try to run it on GPU using cuda, it seems I can no longer get the gradient.</p>
<p>In the code below, if device = ‘cpu’, then I can get the sparse and dense gradient for a and b respectively. If device = ‘cuda’, then I get the error that<br>
<code>&lt;ipython-input-104-da255cd7360a&gt;:1: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more information.</code><br>
This applies for both torch.sparse.mm or torch.spmm method.</p>
<p>Is it the case that gradient using sparse matrix is in general not available on cuda?</p>
<pre><code class="lang-auto">import torch
device='cpu'
# sparse
a = torch.randint(0, 2, size=(2, 3)).float().to_sparse().requires_grad_(True).to(device)
a
# dense
b = torch.randn(3, 2, requires_grad=True).to(device)
b
y = torch.sparse.mm(a, b)
# y = torch.spmm(a, b)
y
y.sum().backward()
a.grad
b.grad
</code></pre>

ptrblck: <aside class="quote no-group" data-username="dwg" data-post="1" data-topic="125781">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/a183cd/48.png" class="avatar"> dwg:</div>
<blockquote>
<p>Is it the case that gradient using sparse matrix is in general not available on cuda?</p>
</blockquote>
</aside>
<p>No, that’s not the case.<br>
As the warning explains, you are trying to access the <code>.grad</code> attribute of a non-leaf tensor, which won’t work.<br>
The reason you are seeing this warning only when using the GPU is because the <code>to(device='cpu')</code> operation will be a no-op on the CPU and a valid operation on the GPU.<br>
You can check it additionally via <code>print(a.is_leaf)</code> and <code>print(b.is_leaf)</code>, which will show <code>True</code> on the CPU and <code>False</code> on the GPU.<br>
This code snippet should work:</p>
<pre data-code-wrap="python"><code class="lang-python">device='cuda'
# sparse
a = torch.randint(0, 2, size=(2, 3)).float().to_sparse().to(device).requires_grad_(True)
a
# dense
b = torch.randn(3, 2, device=device, requires_grad=True)
b
y = torch.sparse.mm(a, b)
y
y.sum().backward()
a.grad
b.grad
</code></pre>

dwg: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="125781">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code>to(device='cpu')</code> operation will be a no-op on the CPU and a valid operation on the GPU.<br>
You can check it additionally via <code>print(a.is_leaf)</code> and <code>print(b.is_leaf)</code>, which will show <code>True</code> on the CPU and <code>False</code> on the GPU.</p>
</blockquote>
</aside>
<p>Awesome! Thank you for the insightful explanation Peter!</p>