Fix autograd issue with identity conversions

[alexsamardzic]

Stack from ghstack (oldest at bottom):

• -> Fix autograd issue with identity conversions #92022

cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92022

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e5505c0:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[alexsamardzic]

This fixes an issue with autograd, for case when the target layout is the same as the original layout. Namely, corresponding conversion methods were doing return self that would result in RuntimeError: leaf variable has been moved into the graph interior for backward operation. Here is a snippet to reproduce the issue:

csr = torch.sparse_csr_tensor((0, 1, 2), (0, 1), (1, 1), dtype=torch.float32, requires_grad=True)
csr2 = csr.to_sparse(layout=torch.sparse_csr)
x = torch.ones((2, 1), dtype=torch.float32)
y = torch.matmul(csr2, x)
z = torch.sum(y)
z.backward()
print(csr.grad)

[cpuhrsch]

The additional clone will introduce a memory copy. Maybe using alias or such can help. @albanD - do you have time to take a look?

[pearu]

I have always wondered why sparse compressed conversions return a copy while sparse or strided conversions return self when layouts do not change. For instance:

>>> a=torch.tensor([[1, 2], [3, 4]])
>>> a.to_dense() is a
True

that makes me assume that csr->csr, csc->csc, ... conversions should also be able to return self (provided that all attributes such as dense_dim, layout, dtype, and device, are unchanged).

Is there a reason for this discrepancy (returning self or self.clone()) between strided and sparse layouts (when tensor attributes are unchanged)?

[nikitaved]

I think it used to break autograd, but that might have changed, cc @albanD .

[albanD]

At a high level the idea here is that autograd elementary ops must respect the aliasing from their schema.
And op is an autograd elementary op iif its key is NOT CompositeImplicitAutograd.

btw this is why ops like reshape must remain CompositeImplicitAutograd. Because they sometimes return a view and sometimes not.

I have always wondered why sparse compressed conversions return a copy while sparse or strided conversions return self when layouts do not change.

t.to_dense() is CompositeImplicitAutograd so you can do either. Returning self here allows you to be faster so we do it.

Any op that is CompositeExplicitAutograd or has special CPU/Sparse dispatch (and thus an autograd formula) cannot do that and must either always be a view or never.

The special note here is that sparse layouts break that rule today for view ops where they return non-views. So when autograd is involved you can get arbitrarily wrong results indeed!

[pearu]

Notice that the following works:

csr = torch.sparse_csr_tensor((0, 1, 2), (0, 1), (1, 1), dtype=torch.float32, requires_grad=True)
csr2 = csr.to_sparse(layout=torch.sparse_csr).detach().requires_grad_(True)
x = torch.ones((2, 1), dtype=torch.float32)
y = torch.matmul(csr2, x)
z = torch.sum(y)
z.backward()
print(csr2.grad)   # UPDATED

What is the autograd interpretation of using .detach().requires_grad_(True) instead of .clone() in .to_sparse()? Would this be a meaningful fix to the autograd issue that this PR targets?

[nikitaved]

csr = torch.sparse_csr_tensor((0, 1, 2), (0, 1), (1, 1), dtype=torch.float32, requires_grad=True)
csr2 = csr.to_sparse(layout=torch.sparse_csr).detach().requires_grad_(True)
x = torch.ones((2, 1), dtype=torch.float32)
y = torch.matmul(csr2, x)
z = torch.sum(y)
z.backward()
print(csr.grad)

It does not seem to work for me, prints None.

[pearu]

It does not seem to work for me, prints None.

Yes, returning None from csr.grad makes sense. I have updated the example, the print statement should read print(csr2.grad).

This actually answers my question in a way: replacing clone with detach() is not meaningful for autograd.

To keep to_sparse efficient for non-autograd context, an option is to introduce copy=True to to_sparse method. Does this makes sense?

[alexsamardzic]

So this PR is basically replacing return self; with return self.clone(); on couple places in order to have the gradient calculations correct. How about following instead then?

return self.requires_grad() ? self.clone() : self;

[amjames]

Tests are green and this looks good from my end. Pending @pearu indicating his concerns have been addressed we can land.

Also regarding the issue of the sparse scalar concept that @pearu pointed out I would say it is out of scope for this work and can be discussed elsewhere at a time when we have a motivation for such a feature.

[pearu]

I have a number of nits and questions. Overall, looks good to me. Thanks, @alexsamardzic!

[pearu]

Btw, here and below, will need return self.alias() but not in this PR as it is handled by #103810

[alexsamardzic]

Great! I should remember to make the change when this PR of yours landed.

[pearu]

Should these dates be updated, say, increment by half a year?

[alexsamardzic]

Changed.

[alexsamardzic]

One more comment here: I've modeled this PR according to the to_dense()/_to_dense() separation. Here, in native_functions.yaml file to_dense is written so that to_dense.out is not generated, and this is the reason I've removed .out versions for to_sparse*() methods. But it could be a backward incompatible change in case someone used this. Please advise if it would be preferable to keep these.

[pearu]

@alexsamardzic any idea why CI tests fail? It looks like to_sparse aten implementation is never called...

[alexsamardzic]

@alexsamardzic any idea why CI tests fail? It looks like to_sparse aten implementation is never called...

Typo 😳

[pearu]

LGTM. Thanks, @alexsamardzic!

[alexsamardzic]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here