Implement the general log-func-exp trick

[lkct]

All the log-*-exp use cases can be generalized to this.

For simple log-sum-exp where we only do a sum, use torch.logsumexp instead.