Simplify logstatsexp reductions: length-aware iterators and direct normalization

docs/src/index.md

@@ -38,4 +38,10 @@ loglogistic
 logitexp
 log1mlogistic
 logit1mexp
+logmeanexp
+logmeanexp!
+logvarexp
+logvarexp!
+logstdexp
+logstdexp!
 ```

ext/LogExpFunctionsChainRulesCoreExt.jl

@@ -151,6 +151,57 @@ function ChainRulesCore.rrule(::typeof(logsumexp), x::AbstractArray{<:Real}; dim
     return Ω, logsumexp_pullback
 end
 
+function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims=:)
+    Ω = logmeanexp(x; dims)
+    n = length(x) ÷ length(Ω)
+    ΔΩ = sum(exp.(x .- Ω) .* Δx; dims) ./ n
+    return Ω, ΔΩ
+end
+function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims=:)
+    Ω = logmeanexp(x; dims)
+    n = length(x) ÷ length(Ω)
+    return Ω, _∂x_pullback(exp.(x .- Ω) ./ n, x)
+end
+
+function ChainRulesCore.frule((_, Δx), ::typeof(logvarexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true)
+    logmean = logmeanexp(x; dims)
+    Ω = logvarexp(x; dims, corrected, logmean)
+    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) .* Δx; dims)
+    return Ω, ΔΩ
+end
+function ChainRulesCore.rrule(::typeof(logvarexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true)
+    logmean = logmeanexp(x; dims)
+    Ω = logvarexp(x; dims, corrected, logmean)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims), x)
+end
+function ChainRulesCore.frule((_, Δx), ::typeof(logstdexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true)
+    logmean = logmeanexp(x; dims)
+    Ω = logstdexp(x; dims, corrected, logmean)
+    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) ./ 2 .* Δx; dims)
+    return Ω, ΔΩ
+end
+function ChainRulesCore.rrule(::typeof(logstdexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true)
+    logmean = logmeanexp(x; dims)
+    Ω = logstdexp(x; dims, corrected, logmean)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims) / 2, x)
+end
+function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims)
+    d = x .- logmean
+    e = expm1.(d)
+    return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims)
+end
+function _∂x_pullback(∂x, x)
+    project_x = ChainRulesCore.ProjectTo(x)
+    function pullback(Ω̄)
+        x̄ = ChainRulesCore.InplaceableThunk(
+            Δ -> Δ .+= Ω̄ .* ∂x,
+            ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)),
+        )
+        return ChainRulesCore.NoTangent(), x̄
+    end
+    return pullback
+end
+
 # no rules for mutating functions currently:
 # https://juliadiff.org/ChainRulesCore.jl/stable/writing_good_rules.html#Which-functions-need-rules?
 function ChainRulesCore.frule((_, Δx), ::typeof(softmax), x::AbstractArray{<:Real}; dims=:)

src/LogExpFunctions.jl

@@ -9,10 +9,12 @@ 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, logabssinh, logabstanh, cloglog, cexpexp,
+    logmeanexp, logmeanexp!, logvarexp, logvarexp!, logstdexp, logstdexp!,
     loglogistic, logitexp, log1mlogistic, logit1mexp
 
 include("basicfuns.jl")
 include("logsumexp.jl")
+include("logstatsexp.jl")
 
 if !isdefined(Base, :get_extension)
     include("../ext/LogExpFunctionsChainRulesCoreExt.jl")

src/logstatsexp.jl

@@ -0,0 +1,106 @@
+# Numerically stable `log`-of-statistics-of-`exp` reductions:
+# logmeanexp, logvarexp, logstdexp.
+
+"""
+$(SIGNATURES)
+
+Compute `log(mean(exp, X))` in a numerically stable way.
+
+`X` should be an iterator of numbers. For an array, `dims` selects the dimensions to reduce
+over, returning `log.(mean(exp.(X); dims))`.
+
+See also [`logmeanexp!`](@ref).
+"""
+function logmeanexp(X)
+    next = iterate(X)
+    isnothing(next) && throw(ArgumentError("reducing over an empty collection is not allowed"))
+    x, state = next
+    acc = x
+    count = 1
+    while true
+        next = iterate(X, state)
+        isnothing(next) && break
+        x, state = next
+        acc = _logsumexp_onepass_op(acc, x)
+        count += 1
+    end
+    lse = _logsumexp_onepass_result(acc)
+    return oftype(lse, lse - log(count))
+end
+logmeanexp(X::AbstractArray{<:Number}; dims=:) = _logmeanexp(X, dims)
+function _logmeanexp(X::AbstractArray{<:Number}, ::Colon)
+    out = logsumexp(X)
+    return oftype(out, out - log(length(X)))
+end
+function _logmeanexp(X::AbstractArray{<:Number}, dims)
+    out = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims))
+    return logmeanexp!(out, X)
+end
+
+"""
+$(SIGNATURES)
+
+Compute [`logmeanexp`](@ref) of `X` over the singleton dimensions of `out`, and write the
+result to `out`.
+"""
+function logmeanexp!(out::AbstractArray, X::AbstractArray{<:Number})
+    logsumexp!(out, X)
+    return out .-= log(length(X) / length(out))
+end
+
+
+"""
+$(SIGNATURES)
+
+Compute `log(var(exp.(X); corrected))` in a numerically stable way.
+
+`X` should be an array of real numbers; `dims` selects the dimensions to reduce over,
+returning `log.(var(exp.(X); dims, corrected))`. A precomputed `logmeanexp(X; dims)` can be
+passed as `logmean` to avoid recomputing it.
+
+See also [`logvarexp!`](@ref).
+"""
+logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims)) = _logvarexp(X, dims, corrected, logmean)
+function _logvarexp(X::AbstractArray{<:Real}, ::Colon, corrected::Bool, logmean)
+    out = logsumexp(2 .* logsubexp.(X, logmean))
+    return oftype(out, out - log(max(0, length(X) - corrected)))
+end
+function _logvarexp(X::AbstractArray{<:Real}, dims, corrected::Bool, logmean)
+    out = logsumexp(2 .* logsubexp.(X, logmean); dims)
+    return out .-= log(max(0, length(X) / length(out) - corrected))
+end
+
+
+"""
+$(SIGNATURES)
+
+Compute [`logvarexp`](@ref) of `X` over the singleton dimensions of `out`, and write the
+result to `out`.
+"""
+function logvarexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; corrected::Bool=true)
+    logmeanexp!(out, X)
+    logsumexp!(out, 2 .* logsubexp.(X, out))
+    return out .-= log(max(0, length(X) / length(out) - corrected))
+end
+
+
+"""
+$(SIGNATURES)
+
+Compute `log(std(exp.(X); dims, corrected))` in a numerically stable way.
+
+Keyword arguments are as in [`logvarexp`](@ref).
+"""
+logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims)) = logvarexp(X; dims, corrected, logmean) / 2
+
+
+"""
+$(SIGNATURES)
+
+Compute [`logstdexp`](@ref) of `X` over the singleton dimensions of `out`, and write the
+result to `out`.
+"""
+function logstdexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; corrected::Bool=true)
+    logvarexp!(out, X; corrected)
+    return out ./= 2
+end

test/Project.toml

@@ -10,6 +10,7 @@ JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/chainrules.jl

@@ -113,6 +113,18 @@
         test_rrule(logsumexp, x; fkwargs=(dims=dims,))
     end
 
+    for x in (randn(10), randn(10, 8)), dims in (:, 1, 1:2, 2)
+        dims isa Colon || all(d <= ndims(x) for d in dims) || continue
+        test_frule(logmeanexp, x; fkwargs=(dims=dims,))
+        test_rrule(logmeanexp, x; fkwargs=(dims=dims,))
+        for corrected in (true, false)
+            test_frule(logvarexp, x; fkwargs=(dims=dims, corrected=corrected))
+            test_rrule(logvarexp, x; fkwargs=(dims=dims, corrected=corrected))
+            test_frule(logstdexp, x; fkwargs=(dims=dims, corrected=corrected))
+            test_rrule(logstdexp, x; fkwargs=(dims=dims, corrected=corrected))
+        end
+    end
+
     for x in (randn(10), randn(10, 8))
         test_frule(softmax, x)
         test_rrule(softmax, x)