Compute the log of a sum of exponentials

def log_sum_exp(self,x, axis):
        '''Compute the log of a sum of exponentials'''
        x_max = np.max(x, axis=axis)
        if axis == 1:
            return x_max + np.log( np.sum(np.exp(x-x_max[:,np.newaxis]), axis=1) )
        else:
            return x_max + np.log( np.sum(np.exp(x-x_max), axis=0) )