Implement scatter_logsumexp, scatter_softmax, scatter_log_softmax

[mallamanis]

(As discussed in #76)

• scatter_logsumexp code copied from Scatter_logaddexp #71 with small modifications.
• scatter_*softmax operations added in torch_scatter.composite subpackage.
• added some documentation.

Closes #76

Open Issues

@rusty1s opening this for discussion. Haven't tested this at all, yet but need to stop for the day. Let me know if you have any comments so far.

Also, there are some design decisions/issues that are not up to me to decide about:

• In scatter_logsumexp the out argument is a bit cumbersome. I've made a choice about handling this, but let me know if you'd like to change this.
• I've removed some arguments in the scatter_*softmax ops. Are you happy with this?
• Naming: do you prefer scatter_logsumexp or scatter_logaddexp. scatter_logsumexp is "compatible" with the respective PyTorch function. logaddexp is compatible with scatter_add.
• No tests yet. What do you want here?
• What is your preferred math notation in documentation? Let me know if you'd like changes.

[rusty1s]

Overall, this looks really good :) Thank you!

1. The out and fill_value seems indeed a bit strange. Personally, I would omit both in the scatter_max call and use them only in the scatter_add call. fill_value should default to 0 instead of torch.finfo(src.dtype).min in the scatter_add call (if not set).
   Edit: Ups, I guess you are right. The default fill_value should be indeed log(0)=torch.finfo(src.dtype).min, but I guess this breaks the scatter_add call.
2. Yes, this is fine. We do not really need the other arguments here.
3. I prefer logsumexp.
4. We could test whether the output matches its dense version coming from PyTorch, e.g. by passing an equally binned index tensor ([0, 0, 0, 1, 1, 1, 2, 2, 2]).
5. Math notation seems fine to me.

I wonder if it is really beneficial to compute scatter_softmax with the help of scatter_log_softmax. What's the advantage here? I believe subtracting by the max value should be sufficient (as it is done in the PyTorch Geometric implementation).

[codecov-io]

Codecov Report

Merging #77 into master will not change coverage.
The diff coverage is 100%.

@@          Coverage Diff          @@
##           master    #77   +/-   ##
=====================================
  Coverage     100%   100%           
=====================================
  Files          10     13    +3     
  Lines         126    159   +33     
=====================================
+ Hits          126    159   +33

Impacted Files	Coverage Δ
torch_scatter/logsumexp.py	100% <100%> (ø)
torch_scatter/composite/softmax.py	100% <100%> (ø)
torch_scatter/composite/__init__.py	100% <100%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 78a5549...62c6122. Read the comment docs.

[mallamanis]

Some progress for today:

• scatter_*softmax but correctly implemented.
  • The scatter_log_softmax had an implementation bug. I've now fixed this. It now has a standalone implementation, which removes one redundant computation.
  • The scatter_softmax is implemented as you correctly suggested.
  • The fill_value, out in scatter_logsumexp should now be consistent with the rest of the scatter_* functions.
  • Inlined the utility _scatter_logsumexp since it's now used only once.
• Added rudimentary tests. As always, more can be added.
  • I tried to be consistent with the style of your tests, but feel free to edit them/let me know what you'd like to change.
  • Importantly, the dim, out are not tested anywhere yet. I won't be able to further work on this today/tomorrow, but feel free to add more/let me know what you want to see in there.
• Addressed CI/flake8 code style issues.

[rusty1s]

Thanks for the update. I will have a last look and will merge afterwards. I am a bit stressed due to ICLR rebuttal, but I will merge as soon as possible. Sorry for the delay.

[rusty1s]

Merged and released in torch-scatter==1.4.0. Thanks again :)

[mallamanis]

Thanks for accepting this PR and your help. Best wishes for ICLR.