diff --git a/docs/src/index.md b/docs/src/index.md index 9454e0c1..907b25a8 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -38,4 +38,10 @@ loglogistic logitexp log1mlogistic logit1mexp +logmeanexp +logmeanexp! +logvarexp +logvarexp! +logstdexp +logstdexp! ``` diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 28fe8fe7..90344527 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/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=:) diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index cdb917b0..f9c20f09 100644 --- a/src/LogExpFunctions.jl +++ b/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") diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl new file mode 100644 index 00000000..f8fca575 --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index f94caf5f..03308b32 100644 --- a/test/Project.toml +++ b/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] diff --git a/test/chainrules.jl b/test/chainrules.jl index c19ae936..597b8eb7 100644 --- a/test/chainrules.jl +++ b/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) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl new file mode 100644 index 00000000..ec2b8dad --- /dev/null +++ b/test/logstatsexp.jl @@ -0,0 +1,227 @@ +using Test: @test, @test_throws, @testset, @inferred +using Statistics: mean, std, var +using OffsetArrays: OffsetArray +using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logstdexp!, logvarexp, logvarexp! + +# Count heap allocations of `f(x)` after warming it up +allocations(f, x) = (f(x); @allocated f(x)) + +# Single-use iterator that advertises a length: iterating consumes it, so a second pass +# (or `isempty` probe) would skip elements. Exercises `logmeanexp`'s one-pass path. +mutable struct DrainOnce{T} + data::Vector{T} + pos::Int +end +DrainOnce(v) = DrainOnce(collect(v), 0) +Base.IteratorSize(::Type{<:DrainOnce}) = Base.HasLength() +Base.length(d::DrainOnce) = length(d.data) - d.pos +Base.eltype(::Type{DrainOnce{T}}) where {T} = T +function Base.iterate(d::DrainOnce, _=nothing) + d.pos < length(d.data) || return nothing + d.pos += 1 + return d.data[d.pos], nothing +end + +@testset "logmeanexp, logvarexp, logstdexp arrays" begin + for T in (Float32, Float64) + X = randn(T, 5, 3, 2) + for dims in (2, (1, 2), :) + @test logmeanexp(X; dims) ≈ log.(mean(exp.(X); dims)) + for corrected in (true, false) + @test logvarexp(X; dims, corrected) ≈ + log.(var(exp.(X); dims, corrected)) + @test logstdexp(X; dims, corrected) ≈ + log.(std(exp.(X); dims, corrected)) + end + end + @test @inferred(logmeanexp(X)) ≈ log(mean(exp, X)) + @test @inferred(logvarexp(X)) ≈ log(var(exp.(X))) + @test @inferred(logstdexp(X)) ≈ log(std(exp.(X))) + end +end + +@testset "logmeanexp iterators, empty reductions" begin + x = randn(Float32, 20) + xt = Tuple(x) + xg = (v for v in x) + xf = Iterators.filter(_ -> true, x) + + # `logmeanexp` is single-pass, so it still accepts arbitrary iterators + @test @inferred(logmeanexp(xt)) ≈ log(mean(exp, xt)) + @test logmeanexp(xg) ≈ log(mean(exp, x)) + @test @inferred(logmeanexp(xf)) ≈ log(mean(exp, x)) + @test logmeanexp(Iterators.Stateful(x)) ≈ log(mean(exp, x)) + + # empty `Tuple` has unknown eltype, so the reduction errors; empty arrays give NaN + @test_throws ArgumentError logmeanexp(()) + @test isnan(logmeanexp(Float64[])) + @test isnan(logvarexp(Float64[])) +end + +@testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin + X = randn(Float32, 5, 3, 2) + + for dims in (1, (2, 3)) + @test eltype(@inferred(logmeanexp(X; dims))) == Float32 + @test eltype(@inferred(logvarexp(X; dims))) == Float32 + @test eltype(@inferred(logstdexp(X; dims))) == Float32 + end + + @test typeof(@inferred(logmeanexp(X; dims=:))) == Float32 + @test typeof(@inferred(logvarexp(X; dims=:))) == Float32 + @test typeof(@inferred(logstdexp(X; dims=:))) == Float32 + + X1 = reshape(randn(Float64, 8), 1, 8) + @test all(isnan, logvarexp(X1; dims=1, corrected=true)) + @test all(isnan, logstdexp(X1; dims=1, corrected=true)) + Xsingleton = fill(0.0f0, 1, 1, 1) + @test isnan(logvarexp(Xsingleton; dims=:, corrected=true)) + @test isnan(logstdexp(Xsingleton; dims=:, corrected=true)) +end + +@testset "return types" begin + # The result element type must follow `float(eltype(input))` + for (Tin, Tout) in ((Float32, Float32), (Float64, Float64), (Int, Float64)) + X = Tin <: Integer ? rand(Tin(1):Tin(5), 5, 3, 2) : randn(Tin, 5, 3, 2) + + # whole-array (`dims=:`) reductions: scalar of type `Tout` + @test typeof(@inferred(logmeanexp(X))) == Tout + @test typeof(@inferred(logmeanexp(X; dims=:))) == Tout + @test typeof(@inferred(logvarexp(X))) == Tout + @test typeof(@inferred(logvarexp(X; dims=:))) == Tout + @test typeof(@inferred(logstdexp(X))) == Tout + @test typeof(@inferred(logstdexp(X; dims=:))) == Tout + + # `dims` reductions: array with eltype `Tout` + for dims in (1, 2, (1, 2)) + @test eltype(@inferred(logmeanexp(X; dims))) == Tout + @test eltype(@inferred(logvarexp(X; dims))) == Tout + @test eltype(@inferred(logstdexp(X; dims))) == Tout + end + + # single-pass iterator path + @test typeof(@inferred(logmeanexp(Tuple(vec(X))))) == Tout + end + + # in-place reductions preserve `out`'s eltype, incl. an Int input widened into a float `out`. + Xi = rand(1:5, 6, 4) + @test eltype(logmeanexp!(Matrix{Float64}(undef, 1, 4), Xi)) == Float64 + @test eltype(logvarexp!(Matrix{Float64}(undef, 1, 4), Xi)) == Float64 + @test eltype(logstdexp!(Matrix{Float64}(undef, 1, 4), Xi)) == Float64 + Xf = randn(Float32, 6, 4) + @test eltype(logmeanexp!(Matrix{Float32}(undef, 1, 4), Xf)) == Float32 + @test eltype(logvarexp!(Matrix{Float32}(undef, 1, 4), Xf)) == Float32 + @test eltype(logstdexp!(Matrix{Float32}(undef, 1, 4), Xf)) == Float32 +end + +# Regression tests for review-found bugs, one block per bug class. +@testset "edge-case regressions" begin + # Non-1-based axes (OffsetArrays) must match the 1-based result; a fused lazy + # reduction silently returned wrong values here. + base = randn(4, 3) + oa = OffsetArray(base, -1, -1) + for dims in (1, 2, :) + @test collect(logmeanexp(oa; dims)) ≈ collect(logmeanexp(base; dims)) + for corrected in (true, false) + @test collect(logvarexp(oa; dims, corrected)) ≈ + collect(logvarexp(base; dims, corrected)) + @test collect(logstdexp(oa; dims, corrected)) ≈ + collect(logstdexp(base; dims, corrected)) + end + end + + # Abstract eltype: `dims` variance must work and match a concrete copy, not throw MethodError. + Xabstract = Real[1.0 2.0 3.0; 4.0 5.0 6.0] + Xconcrete = Float64.(Xabstract) + @test logvarexp(Xabstract; dims=1) ≈ logvarexp(Xconcrete; dims=1) + @test logvarexp(Xabstract) ≈ logvarexp(Xconcrete) + @test logvarexp(Xabstract; dims=2) ≈ logvarexp(Xconcrete; dims=2) + + # Empty along the reduced dim: NaN, not an error (matching `Statistics.var`), for any `corrected`. + Eredux = Matrix{Float64}(undef, 0, 3) + @test all(isnan, logmeanexp(Eredux; dims=1)) + for corrected in (true, false) + @test all(isnan, logvarexp(Eredux; dims=1, corrected)) + @test all(isnan, logstdexp(Eredux; dims=1, corrected)) + end + + # Empty along a non-reduced dim: empty array of the reduced shape, no DivideError. + Eother = Matrix{Float64}(undef, 3, 0) + @test size(logmeanexp(Eother; dims=1)) == (1, 0) + @test size(logvarexp(Eother; dims=1)) == (1, 0) + @test size(logstdexp(Eother; dims=1)) == (1, 0) + + # Single-use iterator with a length: must be traversed exactly once, not re-consumed. + data = randn(7) + @test logmeanexp(DrainOnce(data)) ≈ log(mean(exp, data)) + + # Complex arrays are rejected (no `<:Real` method) on every variance/std path. + C = ComplexF64[1 2; 3 4] + @test_throws Exception logvarexp(C) + @test_throws Exception logvarexp(C; dims=1) + @test_throws Exception logstdexp(C; dims=2) +end + +@testset "in-place logmeanexp!/logvarexp!/logstdexp!" begin + for T in (Float32, Float64) + X = randn(T, 6, 4) + for A in (X, OffsetArray(X, -2, -1)), dims in (1, 2, (1, 2)) + out = similar(A, T, Base.reduced_indices(axes(A), dims)) + @test logmeanexp!(out, A) === out + @test out ≈ logmeanexp(A; dims) + for corrected in (true, false) + outv = similar(A, T, Base.reduced_indices(axes(A), dims)) + @test logvarexp!(outv, A; corrected) === outv + @test outv ≈ logvarexp(A; dims, corrected) + outs = similar(A, T, Base.reduced_indices(axes(A), dims)) + @test logstdexp!(outs, A; corrected) === outs + @test outs ≈ logstdexp(A; dims, corrected) + end + end + end + Xabstract = Real[1.0 2.0 3.0; 4.0 5.0 6.0] + @test logvarexp!(Matrix{Float64}(undef, 1, 3), Xabstract) ≈ logvarexp(Float64.(Xabstract); dims=1) + @test logstdexp!(Matrix{Float64}(undef, 1, 3), Xabstract) ≈ logstdexp(Float64.(Xabstract); dims=1) + @test_throws Exception logvarexp!(Matrix{ComplexF64}(undef, 1, 2), ComplexF64[1 2; 3 4]) + @test_throws Exception logstdexp!(Matrix{ComplexF64}(undef, 1, 2), ComplexF64[1 2; 3 4]) +end + +@testset "allocations" begin + # `logmeanexp` builds no O(n) temporary, so allocations don't grow with input size. + # (`logvarexp`/`logstdexp` do allocate one, so they aren't checked here.) + for T in (Float32, Float64) + @test allocations(logmeanexp, randn(T, 10_000)) == allocations(logmeanexp, randn(T, 100)) + @test allocations(logmeanexp, randn(T, 10_000)) == 0 + g = X -> logmeanexp(X; dims=1) + @test allocations(g, randn(T, 10_000, 3)) == allocations(g, randn(T, 100, 3)) + out = Matrix{T}(undef, 1, 3) + @test allocations(X -> logmeanexp!(out, X), randn(T, 10_000, 3)) == + allocations(X -> logmeanexp!(out, X), randn(T, 100, 3)) + end +end + +@testset "numerical robustness" begin + # Compare against BigFloat references on hard cases: tight clusters (var ≪ mean², where + # a naive second-moment formula cancels) and values where exp over-/under-flows Float64. + setprecision(BigFloat, 256) do + refmean(x) = Float64(log(mean(exp.(big.(x))))) + refvar(x; corrected=true) = Float64(log(var(exp.(big.(x)); corrected))) + cases = ( + 1.0 .+ 1e-3 .* randn(500), # tight cluster + 1.0 .+ 1e-6 .* randn(500), # very tight: var ≪ mean² + 700.0 .+ randn(200), # exp overflows Float64 (>~709) + -700.0 .+ randn(200), # exp underflows to 0 + 1000.0 .* randn(500), # huge dynamic range + [3.0, 3.0 + 1e-10], # nearly-equal pair + [0.0, 1e-6], + ) + for x in cases + # atol covers near-zero results, where relative error is meaningless + @test logmeanexp(x) ≈ refmean(x) rtol = 1e-9 atol = 1e-10 + for corrected in (true, false) + @test logvarexp(x; corrected) ≈ refvar(x; corrected) rtol = 1e-8 atol = 1e-9 + @test logstdexp(x; corrected) ≈ refvar(x; corrected) / 2 rtol = 1e-8 atol = 1e-9 + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 5fcc6ad8..3c84f957 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,6 +19,7 @@ include("basicfuns.jl") include("chainrules.jl") include("inverse.jl") include("with_logabsdet_jacobian.jl") +include("logstatsexp.jl") # QA import JET