import numpy def logsumexp(x, dim=-1): """Compute numerically stable log(sum(exp(x))). Use second argument to specify along which dimension to sum. If -1 (default), logsumexp is computed along the last dimension. """ if len(x.shape) < 2: #got only one dimension? xmax = x.max() return xmax + numpy.log(numpy.sum(numpy.exp(x-xmax))) else: if dim != -1: x = x.transpose(range(dim) + range(dim+1, len(x.shape)) + [dim]) lastdim = len(x.shape)-1 xmax = x.max(lastdim) return xmax + numpy.log(numpy.sum(numpy.exp(x-xmax[...,newaxis]), lastdim))