Add sparse COO tensor support to torch.sum(dim=..., keepdim=...)

[pearu]

Fixes #92757, #86232

Stack from ghstack (oldest at bottom):

• -> Add sparse COO tensor support to torch.sum(dim=..., keepdim=...) #92979

cc @alexsamardzic @nikitaved @cpuhrsch @amjames @bhosmer

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 44bc03a:
💚 Looks good so far! There are no failures yet. 💚

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

[pearu]

@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

[vadimkantorov]

should these unsqueeze start from the last dimension instead of the first ones?

[pearu]

no, because dim_mask and the dim loop correspond to the self shape and unsqueeze inserts (read: restores) the removed dimensions at the position that is in the reference of self, not of result, shape.

Consider the following example:

>>> a = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
>>> mask_dim = [True, True, False]
>>> dim = tuple(i for i in range(len(mask_dim)) if mask_dim[i])
>>> dim
(0, 1)
>>> # Using algorithm in this PR:
>>> result = a.sum(dim=dims, keepdim=False)
>>> for i in range(a.ndim):
...     result = result.unsqueeze(i) if mask_dim[i] else result
... 
>>> result
tensor([[[16, 20]]])
>>> # Using unsqueeze starting from the last dimension:
>>> result = a.sum(dim=dim, keepdim=False)
>>> for i in reversed(range(a.ndim)):
...     result = result.unsqueeze(i) if mask_dim[i] else result
... 
>>> result
tensor([[[16],
         [20]]])
>>> # The expected result:
>>> a.sum(dim=dim, keepdim=True)
tensor([[[16, 20]]])

As you can see, the algorithm implemented in this PR gives the expected result while reversing the order of unsqueeze not.