Add PyTorch SparseTensor support for MessagePassing
#5944
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
68 tasks
|
Testing failed due to numerical instability. It seems we should use |
Codecov Report
@@ Coverage Diff @@
## master #5944 +/- ##
==========================================
- Coverage 86.36% 84.54% -1.82%
==========================================
Files 361 361
Lines 19849 19872 +23
==========================================
- Hits 17142 16801 -341
- Misses 2707 3071 +364
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
|
This looks pretty good. Is it ready to review? |
|
Yes. Please take a look at your convenience :) |
rusty1s
approved these changes
Nov 15, 2022
|
Sure :) Will also add it to the roadmap. |
jjpietrak
pushed a commit
to jjpietrak/pytorch_geometric
that referenced
this pull request
Nov 16, 2022
This PR adds the PyTorch SparseTensor support for the base layer `MessagePassing`. There are some points to be confirmed (as marked with TODO): + ~~In `__collect__`: Since `adj._values()` returns a detached tensor, should we use a coalesced matrix instead (e.g., `adj.coalesce().values()`)? This is for the case of computing sparse gradients of `adj`.~~ (Solved) + In `__collect__`: Should we store the `ptr` for PyTorch SparseTensor when fused aggregation is not available? + In `__lift__`: Should we use `gather_csr` for PyTorch SparseTensor? Also, `torch.jit.script` is not available for PyTorch SparseTensor. Will figure it out soon. Co-authored-by: rusty1s <[email protected]>
jjpietrak
pushed a commit
to jjpietrak/pytorch_geometric
that referenced
this pull request
Nov 25, 2022
This PR adds the PyTorch SparseTensor support for the base layer `MessagePassing`. There are some points to be confirmed (as marked with TODO): + ~~In `__collect__`: Since `adj._values()` returns a detached tensor, should we use a coalesced matrix instead (e.g., `adj.coalesce().values()`)? This is for the case of computing sparse gradients of `adj`.~~ (Solved) + In `__collect__`: Should we store the `ptr` for PyTorch SparseTensor when fused aggregation is not available? + In `__lift__`: Should we use `gather_csr` for PyTorch SparseTensor? Also, `torch.jit.script` is not available for PyTorch SparseTensor. Will figure it out soon. Co-authored-by: rusty1s <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds the PyTorch SparseTensor support for the base layer
MessagePassing. There are some points to be confirmed (as marked with TODO):In(Solved)__collect__: Sinceadj._values()returns a detached tensor, should we use a coalesced matrix instead (e.g.,adj.coalesce().values())? This is for the case of computing sparse gradients ofadj.__collect__: Should we store theptrfor PyTorch SparseTensor when fused aggregation is not available?__lift__: Should we usegather_csrfor PyTorch SparseTensor?Also,
torch.jit.scriptis not available for PyTorch SparseTensor. Will figure it out soon.