Add torch.logcumsumexp

[kshitij12345]

Creating new PR as I am unable to push to @pandeykartikey 's branch as I don't have the permissions.

Closes #26411

Based on #32876 Thanks @pandeykartikey for starting this out.

Have addressed the comments.

@anjali411 @agadetsky @albanD

[dr-ci]

💊 CI failures summary and remediations

As of commit c0ee7ab (more details on the Dr. CI page):

---

• 1/3 failures possibly* introduced in this PR
  • 1/1 non-CircleCI failure(s)
• 1/3 tentatively recognized as flaky ❄️
  • Click here to rerun these jobs
• 1/3 broken upstream at merge base f3b5c22 since May 19

---

❄️ 1 failure tentatively classified as flaky

but reruns have not yet been triggered to confirm:

pytorch_libtorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_test (1/1)

Step: "Set Up CI Environment After attach_workspace" (full log | diagnosis details | 🔁 rerun) ❄️

E: Failed to fetch https://nvidia.github.io/nvidia-container-runtime/ubuntu16.04/amd64/Packages 404 Not Found

                            100% [Working]                Fetched 5,129 kB in 3s (1,652 kB/s) 
Reading package lists... 99%  Reading package lists... Done  
W: The repository 'https://nvidia.github.io/nvidia-container-runtime/ubuntu16.04/amd64  Release' does not have a Release file. 
N: Data from such a repository can't be authenticated and is therefore potentially dangerous to use. 
N: See apt-secure(8) manpage for repository creation and user configuration details. 
W: An error occurred during the signature verification. The repository is not updated and the previous index files will be used. GPG error: https://cli-assets.heroku.com/apt ./ InRelease: The following signatures couldn't be verified because the public key is not available: NO_PUBKEY 5DC22404A6F9F1CA 
W: The repository 'https://packagecloud.io/circleci/trusty/ubuntu xenial Release' does not have a Release file. 
N: Data from such a repository can't be authenticated and is therefore potentially dangerous to use. 
N: See apt-secure(8) manpage for repository creation and user configuration details. 
W: Failed to fetch https://cli-assets.heroku.com/apt/./InRelease  The following signatures couldn't be verified because the public key is not available: NO_PUBKEY 5DC22404A6F9F1CA 
E: Failed to fetch https://nvidia.github.io/nvidia-container-runtime/ubuntu16.04/amd64/Packages  404  Not Found 
W: Some index files failed to download. They have been ignored, or old ones used instead. 

---

🚧 1 ongoing upstream failure:

These were probably caused by upstream breakages that are not fixed yet:

• pytorch_macos_10_13_py3_test since May 19
  • 🔁 rerun

---

ci.pytorch.org: 1 failed

• Failed: pr/py3.6-clang7-rocmdeb-ubuntu16.04

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker.

See how this bot performed.

This comment has been revised 100 times.

[tridao]

I don't think the current implementation is numerically stable. The original issue #26411 suggests using cummax to reduce numerical error, but that has quadratic complexity (see discussion in #32876).
I still think it's better to implement it the way Tensorflow does it: as a prefix scan of the binary operation log_add_exp.

[kshitij12345]

@anjali411 @albanD Could you please review the PR. Also I don't have the rights to re-run the failed pipeline (whose failure is unrelated to the PR).

@tridao I get the point . But I am not sure that I am quite familiar with the codebase. I saw the reference for cumsum https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/ReduceOpsKernel.cpp
and https://github.com/pytorch/pytorch/blob/master/aten/src/THC/generic/THCTensorMathScan.cu ( CUDA Implementation which I don't think should be used as reference).
Would be great if I can get some pointers.

[tridao]

Agreed, those seem to be where cumsum is implemented.
I imagine one just needs to replace the addition (x + y) of cumsum with the operation log_add_exp(x, y)(implemented in a stable way with max and min, e.g. https://www.tensorflow.org/api_docs/python/tf/math/cumulative_logsumexp).

Re: CUDA implementation in THC: maybe cumsum will eventually be ported from THC to Aten? I also found cummax implementation in Aten (https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/ReduceOpsKernel.cu). One would replace max(x, y) with log_add_exp. The cummax implementation also needs the indices but I don't think logcumsumexp will need those.

[kshitij12345]

@ngimel @albanD Gentle ping :).

[albanD]

Looks good to me.
The perf for the backward might not be great but it is better than nothing.
Feel free to open a new issue if you think the backward should be implemented with a special kernel to make it more efficient. But I think this should be left for a future PR anyway.

[facebook-github-bot]

@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[kshitij12345]

@albanD
Would be great if this lands soon into master as it would potentially unblock #36458 .
Thank You.

[albanD]

The landing is in progress. But there are some flaky internal tests so I had to re-run these...

[facebook-github-bot]

@albanD merged this pull request in 3487744.

[agadetsky]

@kshitij12345 thank you very much, great work!

[agadetsky]

@kshitij12345 seems like documentation is a bit broken. At least computation formula does not appear at the moment at the master doc (https://pytorch.org/docs/master/generated/torch.logcumsumexp.html#torch.logcumsumexp)

[kshitij12345]

@agadetsky Thanks for notifying. Will try to get it fixed.