sumlog

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.3.14"
+version = "0.3.15"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/LogExpFunctions.jl

@@ -11,12 +11,13 @@ import LinearAlgebra
 
 export xlogx, xlogy, xlog1py, xexpx, xexpy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
     softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, logsumexp!, softmax,
-    softmax!, logcosh
+    softmax!, logcosh, sumlog
 
 include("basicfuns.jl")
 include("logsumexp.jl")
 include("chainrules.jl")
 include("inverse.jl")
 include("with_logabsdet_jacobian.jl")
+include("sumlog.jl")
 
 end # module

src/sumlog.jl

@@ -0,0 +1,42 @@
+"""
+$(SIGNATURES)
+
+Compute `sum(log.(X))` with a single `log` evaluation.
+
+This is faster than computing `sum(log.(X))` or even `sum(log, X)`, in
+particular as the size of `X` increases.
+
+This works by representing the `j`th element of `X` as ``x_j = a_j  2^{b_j}``,
+allowing us to write
+```math
+\\sum_j \\log{x_j} = \\log(\\prod_j a_j) + \\log{2} \\sum_j b_j
+```
+Since ``\\log{2}`` is constant, `sumlog` only requires a single `log`
+evaluation.
+"""
+function sumlog(x)
+    T = float(eltype(x))
+    _sumlog(T, values(x))
+end
+
+@inline function _sumlog(::Type{T}, x) where {T<:AbstractFloat}
+    sig, ex = mapreduce(_sumlog_op, x; init=(one(T), zero(exponent(one(T))))) do xj
+        float_xj = float(xj)
+        significand(float_xj), exponent(float_xj) 
+    end
+    return log(sig) + IrrationalConstants.logtwo * ex
+end
+
+@inline function _sumlog_op((sig1, ex1), (sig2, ex2))
+    sig = sig1 * sig2
+    ex = ex1 + ex2
+    # Significands are in the range [1,2), so multiplication will eventually overflow
+    if sig > floatmax(typeof(sig)) / 2
+        ex += exponent(sig)
+        sig = significand(sig)
+    end
+    return sig, ex
+end
+
+# `float(T)` is not always `isa AbstractFloat`, e.g. dual numbers or symbolics
+@inline _sumlog(::Type{T}, x) where {T} = sum(log, x)

test/runtests.jl

@@ -14,3 +14,4 @@ include("basicfuns.jl")
 include("chainrules.jl")
 include("inverse.jl")
 include("with_logabsdet_jacobian.jl")
+include("sumlog.jl")

test/sumlog.jl

@@ -0,0 +1,7 @@
+@testset "sumlog" begin
+    for T in [Int, Float16, Float32, Float64, BigFloat]
+        for x in [10 .* rand(1000), repeat([nextfloat(1.0)], 1000), repeat([prevfloat(2.0)], 1000)]
+            @test (@inferred sumlog(x)) ≈ sum(log, x)
+        end
+    end
+end