Leaving this issue here for higher visibility. Originally reported by @vadimkantorov here: #7254

Also mentioned by @vadimkantorov:

it seems PyTorch now has two copies of logsumexp, second one is in distributions: #8031

Piggybacking off of this issue, torch.distributions.utils.log_sum_exp has strange keepdim behavior (reported here: https://discuss.pytorch.org/t/unexpected-behaviour-with-torch-distributions-utils-log-sum-exp/19594)

torch.distributions.utils.log_sum_exp(torch.ones([10, 100]), keepdim=True).shape
torch.Size([10, 1])
torch.distributions.utils.log_sum_exp(torch.ones([10, 100]), keepdim=False).shape
torch.Size([10, 10])