Although this is stable, the implementation will give incorrect grad when logits=0 (only at this single point) due to torch.abs and torch.clamp operators. I'm not sure how to resolve it. We might reconsider the version pre-#15962.

It turns out that the issue lies at torch.clamp which returns grad at 0 is 1. If it returns 0.5, grad will be correct.

Related #2195 where the same issue happens for binary cross entropy with logits. If we make grad of torch.clamp at 0 is 0.5, the result for grad at 0 of binomal logprob and binary cross entropy with logits (pre #2195) will be correct. @apaszke What do you think?

Nevermind, we have x = x.clamp(min=0) + x.clamp(max=0), the LHS has grad at 0 is 1 while RHS is 2. We can utilize this difference to add a dummy variable (x - x.clamp(min=0) - x.clamp(max=0)) / 2 to x.clamp(min=0) to adjust its grad.

This is ready to review. I think that I have addressed problems which I can think about:

• overflow of log_prob
• discontinuity of grad of log_prob
• correct grad at logits=0

@pytorchbot merge this please

@ezyang merged this pull request in 0719714.