Hi Jethro,
Unfortunately, I can’t find the source code (this was for a toy example that I might have deleted).
However, scipy has an (open source) implementation of the log-sum-exp operation (see below), which is easy to adapt for PyTorch.
Moreover, the idea behind this function is very simple, so if you understand it you might not even need to look at the scipy code. Let me explain.
Log-sum-exp, like its name says, computes the logarithm of a sum of exponentials:

If the values inside the exponentials are all large negative numbers, the sum inside the log vanishes and we get the logarithm of 0, so we’re in trouble. To avoid this, the following simple trick is used:


Find the maximum value among the a_i’s:



Compute b using:



Now, at least one of the values inside the exponentials will be 0, and exp(0) = 1, so we will certainly not get log(0). Thus, we now have a numerically robust implementation of the log-sum-exp operation.
Hope this explanation helps 

  
      github.com
  
  
    scipy/scipy/blob/v0.19.1/scipy/special/_logsumexp.py#L8-L127 16

def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
"""Compute the log of the sum of exponentials of input elements.


Parameters
----------
a : array_like
    Input array.
axis : None or int or tuple of ints, optional
    Axis or axes over which the sum is taken. By default `axis` is None,
    and all elements are summed.


    .. versionadded:: 0.11.0
keepdims : bool, optional
    If this is set to True, the axes which are reduced are left in the
    result as dimensions with size one. With this option, the result
    will broadcast correctly against the original array.


    .. versionadded:: 0.15.0
b : array-like, optional
    Scaling factor for exp(`a`) must be of the same shape as `a` or


  This file has been truncated. show original