From 5b96b87fe807e7e2fd82ff416d5c4f78daa127f6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 18:48:47 +0000 Subject: [PATCH 01/40] Add logmeanexp/logvarexp/logstdexp implementations --- docs/src/index.md | 3 ++ src/LogExpFunctions.jl | 2 + src/logstatsexp.jl | 105 +++++++++++++++++++++++++++++++++++++++++ test/Project.toml | 1 + test/logstatsexp.jl | 34 +++++++++++++ test/runtests.jl | 1 + 6 files changed, 146 insertions(+) create mode 100644 src/logstatsexp.jl create mode 100644 test/logstatsexp.jl diff --git a/docs/src/index.md b/docs/src/index.md index 9454e0c1..903acff5 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -38,4 +38,7 @@ loglogistic logitexp log1mlogistic logit1mexp +logmeanexp +logvarexp +logstdexp ``` diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index cdb917b0..06099f7a 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, logvarexp, 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..2a0b5dbd --- /dev/null +++ b/src/logstatsexp.jl @@ -0,0 +1,105 @@ +""" +$(SIGNATURES) + +Compute `log(mean(exp, X))`. + +`X` should be an iterator of numbers. +The result is computed in a numerically stable way. +""" +function logmeanexp(X) + lse, n = _logsumexp_count(X) + return lse - log(_convert_count(lse, n)) +end + +""" +$(SIGNATURES) + +Compute `log.(mean(exp.(X); dims=dims))`. + +The result is computed in a numerically stable way. +""" +function logmeanexp(X::AbstractArray{<:Number}; dims=:) + R = logsumexp(X; dims=dims) + n = length(X) ÷ length(R) + return _subtract_log_count(R, n) +end + +""" +$(SIGNATURES) + +Compute `log(var(exp, X; corrected=corrected))`. + +`X` should be an iterator of numbers. +The result is computed in a numerically stable way. +""" +function logvarexp(X; corrected::Bool=true, logmean=logmeanexp(X)) + R = logsumexp(x -> 2logsubexp(x, logmean), X) + n = _count_elements(X) + denom = corrected ? n - 1 : n + return R - log(_convert_count(R, denom)) +end + +""" +$(SIGNATURES) + +Compute `log.(var(exp.(X); dims=dims, corrected=corrected))`. + +The result is computed in a numerically stable way. +""" +function logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims)) + R = logsumexp(2logsubexp.(X, logmean); dims=dims) + n = length(X) ÷ length(R) + denom = corrected ? n - 1 : n + return _subtract_log_count(R, denom) +end + +""" +$(SIGNATURES) + +Compute `log(std(exp, X; corrected=corrected))`. + +`X` should be an iterator of numbers. +The result is computed in a numerically stable way. +""" +function logstdexp(X; corrected::Bool=true, logmean=logmeanexp(X)) + return logvarexp(X; corrected=corrected, logmean=logmean) / 2 +end + +""" +$(SIGNATURES) + +Compute `log.(std(exp.(X); dims=dims, corrected=corrected))`. + +The result is computed in a numerically stable way. +""" +function logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims)) + return logvarexp(X; dims=dims, corrected=corrected, logmean=logmean) / 2 +end + +function _logsumexp_count(X) + state = iterate(X) + state === nothing && throw(ArgumentError("reducing over an empty collection is not allowed")) + x, iter_state = state + acc = x + n = 1 + while true + state = iterate(X, iter_state) + state === nothing && break + x, iter_state = state + n += 1 + acc = _logsumexp_onepass_op(acc, x) + end + return _logsumexp_onepass_result(acc), n +end + +_convert_count(x, n::Integer) = convert(typeof(x), n) +_count_elements(X) = _count_elements(X, Base.IteratorSize(typeof(X))) +_count_elements(X, ::Union{Base.HasLength,Base.HasShape}) = length(X) +_count_elements(X, ::Any) = count(_ -> true, X) + +_subtract_log_count(R::Number, n::Integer) = R - log(_convert_count(R, n)) +function _subtract_log_count(R::AbstractArray{<:Number}, n::Integer) + logn = log(convert(eltype(R), n)) + R .-= logn + return R +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/logstatsexp.jl b/test/logstatsexp.jl new file mode 100644 index 00000000..15547f2f --- /dev/null +++ b/test/logstatsexp.jl @@ -0,0 +1,34 @@ +using Test: @test, @testset, @inferred +using Statistics: mean, std, var +using LogExpFunctions: logmeanexp, logstdexp, logvarexp + +@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=dims) ≈ log.(mean(exp.(X); dims=dims)) + for corrected in (true, false) + @test logvarexp(X; dims=dims, corrected=corrected) ≈ + log.(var(exp.(X); dims=dims, corrected=corrected)) + @test logstdexp(X; dims=dims, corrected=corrected) ≈ + log.(std(exp.(X); dims=dims, corrected=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, logvarexp, logstdexp iterators" begin + x = randn(Float32, 20) + xt = Tuple(x) + xg = (v for v in x) + + @test @inferred(logmeanexp(xt)) ≈ log(mean(exp, xt)) + @test logmeanexp(xg) ≈ log(mean(exp, x)) + @test @inferred(logvarexp(xt)) ≈ log(var(exp, xt)) + @test logvarexp(xt; corrected=false) ≈ log(var(exp, xt; corrected=false)) + @test @inferred(logstdexp(xt)) ≈ log(std(exp, xt)) + @test logstdexp(xt; corrected=false) ≈ log(std(exp, xt; corrected=false)) +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 From aa991caad244e00740046cc5b1ece05787d62e2d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 18:55:43 +0000 Subject: [PATCH 02/40] Fix logstats iterator implementation and tests --- src/logstatsexp.jl | 2 +- test/logstatsexp.jl | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 2a0b5dbd..7396ca11 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -33,7 +33,7 @@ Compute `log(var(exp, X; corrected=corrected))`. The result is computed in a numerically stable way. """ function logvarexp(X; corrected::Bool=true, logmean=logmeanexp(X)) - R = logsumexp(x -> 2logsubexp(x, logmean), X) + R = logsumexp((2logsubexp(x, logmean) for x in X)) n = _count_elements(X) denom = corrected ? n - 1 : n return R - log(_convert_count(R, denom)) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 15547f2f..ce68b54a 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -15,8 +15,8 @@ using LogExpFunctions: logmeanexp, logstdexp, logvarexp 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)) + @test @inferred(logvarexp(X)) ≈ log(var(exp.(X))) + @test @inferred(logstdexp(X)) ≈ log(std(exp.(X))) end end @@ -24,11 +24,12 @@ end x = randn(Float32, 20) xt = Tuple(x) xg = (v for v in x) + xe = exp.(x) @test @inferred(logmeanexp(xt)) ≈ log(mean(exp, xt)) @test logmeanexp(xg) ≈ log(mean(exp, x)) - @test @inferred(logvarexp(xt)) ≈ log(var(exp, xt)) - @test logvarexp(xt; corrected=false) ≈ log(var(exp, xt; corrected=false)) - @test @inferred(logstdexp(xt)) ≈ log(std(exp, xt)) - @test logstdexp(xt; corrected=false) ≈ log(std(exp, xt; corrected=false)) + @test @inferred(logvarexp(xt)) ≈ log(var(xe)) + @test logvarexp(xt; corrected=false) ≈ log(var(xe; corrected=false)) + @test @inferred(logstdexp(xt)) ≈ log(std(xe)) + @test logstdexp(xt; corrected=false) ≈ log(std(xe; corrected=false)) end From 3930a786fece6d2d384315b9997670ad91494545 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 19:25:53 +0000 Subject: [PATCH 03/40] Add PR77-focused logstatsexp tests --- test/logstatsexp.jl | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index ce68b54a..0d8d60cb 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -24,12 +24,32 @@ end x = randn(Float32, 20) xt = Tuple(x) xg = (v for v in x) + xf = Iterators.filter(_ -> true, x) xe = exp.(x) @test @inferred(logmeanexp(xt)) ≈ log(mean(exp, xt)) @test logmeanexp(xg) ≈ log(mean(exp, x)) + @test @inferred(logmeanexp(xf)) ≈ log(mean(exp, x)) @test @inferred(logvarexp(xt)) ≈ log(var(xe)) @test logvarexp(xt; corrected=false) ≈ log(var(xe; corrected=false)) + @test @inferred(logvarexp((v for v in x))) ≈ log(var(xe)) + @test logvarexp((v for v in x); corrected=false) ≈ log(var(xe; corrected=false)) @test @inferred(logstdexp(xt)) ≈ log(std(xe)) @test logstdexp(xt; corrected=false) ≈ log(std(xe; corrected=false)) + @test @inferred(logstdexp((v for v in x))) ≈ log(std(xe)) + @test logstdexp((v for v in x); corrected=false) ≈ log(std(xe; corrected=false)) +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=dims))) == Float32 + @test eltype(@inferred(logvarexp(X; dims=dims))) == Float32 + @test eltype(@inferred(logstdexp(X; dims=dims))) == Float32 + end + + @test typeof(@inferred(logmeanexp(X; dims=:))) == Float32 + @test typeof(@inferred(logvarexp(X; dims=:))) == Float32 + @test typeof(@inferred(logstdexp(X; dims=:))) == Float32 end From 0431c847adb99dc76aacfe182f6c19b3420da89d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 19:46:57 +0000 Subject: [PATCH 04/40] Fix PR review issues in logstatsexp --- src/logstatsexp.jl | 57 +++++++++++++++++++++++++++++++++++++-------- test/logstatsexp.jl | 12 ++++++++++ 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 7396ca11..93f2d427 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -20,7 +20,7 @@ The result is computed in a numerically stable way. """ function logmeanexp(X::AbstractArray{<:Number}; dims=:) R = logsumexp(X; dims=dims) - n = length(X) ÷ length(R) + n = _reduced_count(X, R) return _subtract_log_count(R, n) end @@ -32,11 +32,10 @@ Compute `log(var(exp, X; corrected=corrected))`. `X` should be an iterator of numbers. The result is computed in a numerically stable way. """ -function logvarexp(X; corrected::Bool=true, logmean=logmeanexp(X)) - R = logsumexp((2logsubexp(x, logmean) for x in X)) - n = _count_elements(X) +function logvarexp(X; corrected::Bool=true, logmean=nothing) + R, n = isnothing(logmean) ? _logvariance_terms(X) : _logvariance_terms(X, logmean) denom = corrected ? n - 1 : n - return R - log(_convert_count(R, denom)) + return _subtract_log_count(R, denom) end """ @@ -46,9 +45,11 @@ Compute `log.(var(exp.(X); dims=dims, corrected=corrected))`. The result is computed in a numerically stable way. """ -function logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims)) +function logvarexp( + X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims) +) R = logsumexp(2logsubexp.(X, logmean); dims=dims) - n = length(X) ÷ length(R) + n = _reduced_count(X, R) denom = corrected ? n - 1 : n return _subtract_log_count(R, denom) end @@ -61,7 +62,7 @@ Compute `log(std(exp, X; corrected=corrected))`. `X` should be an iterator of numbers. The result is computed in a numerically stable way. """ -function logstdexp(X; corrected::Bool=true, logmean=logmeanexp(X)) +function logstdexp(X; corrected::Bool=true, logmean=nothing) return logvarexp(X; corrected=corrected, logmean=logmean) / 2 end @@ -72,7 +73,9 @@ Compute `log.(std(exp.(X); dims=dims, corrected=corrected))`. The result is computed in a numerically stable way. """ -function logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims)) +function logstdexp( + X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims) +) return logvarexp(X; dims=dims, corrected=corrected, logmean=logmean) / 2 end @@ -97,9 +100,43 @@ _count_elements(X) = _count_elements(X, Base.IteratorSize(typeof(X))) _count_elements(X, ::Union{Base.HasLength,Base.HasShape}) = length(X) _count_elements(X, ::Any) = count(_ -> true, X) -_subtract_log_count(R::Number, n::Integer) = R - log(_convert_count(R, n)) +function _logvariance_terms(X) + lse, lse2, n = _logsumexp2_count(X) + logn = log(_convert_count(lse, n)) + return logsubexp(lse2, 2lse - logn), n +end +function _logvariance_terms(X, logmean) + R, n = _logsumexp_count((2logsubexp(x, logmean) for x in X)) + return R, n +end +function _logsumexp2_count(X) + state = iterate(X) + state === nothing && throw(ArgumentError("reducing over an empty collection is not allowed")) + x, iter_state = state + acc = x + acc2 = 2x + n = 1 + while true + state = iterate(X, iter_state) + state === nothing && break + x, iter_state = state + n += 1 + acc = _logsumexp_onepass_op(acc, x) + acc2 = _logsumexp_onepass_op(acc2, 2x) + end + return _logsumexp_onepass_result(acc), _logsumexp_onepass_result(acc2), n +end + +function _subtract_log_count(R::Number, n::Integer) + n == 0 && return oftype(float(R), NaN) + return R - log(_convert_count(R, n)) +end function _subtract_log_count(R::AbstractArray{<:Number}, n::Integer) + n == 0 && return fill!(R, convert(eltype(R), NaN)) logn = log(convert(eltype(R), n)) R .-= logn return R end + +_reduced_count(X::AbstractArray, R::Number) = length(X) +_reduced_count(X::AbstractArray, R::AbstractArray) = length(X) ÷ length(R) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 0d8d60cb..fa3531b8 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -30,14 +30,19 @@ end @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)) @test @inferred(logvarexp(xt)) ≈ log(var(xe)) @test logvarexp(xt; corrected=false) ≈ log(var(xe; corrected=false)) @test @inferred(logvarexp((v for v in x))) ≈ log(var(xe)) @test logvarexp((v for v in x); corrected=false) ≈ log(var(xe; corrected=false)) + @test logvarexp(Iterators.Stateful(x)) ≈ log(var(xe)) @test @inferred(logstdexp(xt)) ≈ log(std(xe)) @test logstdexp(xt; corrected=false) ≈ log(std(xe; corrected=false)) @test @inferred(logstdexp((v for v in x))) ≈ log(std(xe)) @test logstdexp((v for v in x); corrected=false) ≈ log(std(xe; corrected=false)) + @test logstdexp(Iterators.Stateful(x)) ≈ log(std(xe)) + @test isnan(logvarexp((0.0,))) + @test isnan(logstdexp((0.0,))) end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin @@ -52,4 +57,11 @@ 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 From 6c4283638aeaed92574edc684c457bb7e4158ec7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 20:39:24 +0000 Subject: [PATCH 05/40] Fix unresolved review comments: remove dead _count_elements, update docstrings to say real numbers --- src/logstatsexp.jl | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 93f2d427..5280c1e6 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -29,7 +29,7 @@ $(SIGNATURES) Compute `log(var(exp, X; corrected=corrected))`. -`X` should be an iterator of numbers. +`X` should be an iterator of real numbers. The result is computed in a numerically stable way. """ function logvarexp(X; corrected::Bool=true, logmean=nothing) @@ -59,7 +59,7 @@ $(SIGNATURES) Compute `log(std(exp, X; corrected=corrected))`. -`X` should be an iterator of numbers. +`X` should be an iterator of real numbers. The result is computed in a numerically stable way. """ function logstdexp(X; corrected::Bool=true, logmean=nothing) @@ -96,9 +96,6 @@ function _logsumexp_count(X) end _convert_count(x, n::Integer) = convert(typeof(x), n) -_count_elements(X) = _count_elements(X, Base.IteratorSize(typeof(X))) -_count_elements(X, ::Union{Base.HasLength,Base.HasShape}) = length(X) -_count_elements(X, ::Any) = count(_ -> true, X) function _logvariance_terms(X) lse, lse2, n = _logsumexp2_count(X) From 7f1441241009f6ee22a70e592e8f71405ef26ed5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 20:57:43 +0000 Subject: [PATCH 06/40] Use explicit multiplication in logstatsexp --- src/logstatsexp.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 5280c1e6..330fcf80 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -48,7 +48,7 @@ The result is computed in a numerically stable way. function logvarexp( X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims) ) - R = logsumexp(2logsubexp.(X, logmean); dims=dims) + R = logsumexp(2 * logsubexp.(X, logmean); dims=dims) n = _reduced_count(X, R) denom = corrected ? n - 1 : n return _subtract_log_count(R, denom) @@ -103,7 +103,7 @@ function _logvariance_terms(X) return logsubexp(lse2, 2lse - logn), n end function _logvariance_terms(X, logmean) - R, n = _logsumexp_count((2logsubexp(x, logmean) for x in X)) + R, n = _logsumexp_count((2 * logsubexp(x, logmean) for x in X)) return R, n end function _logsumexp2_count(X) From 28fc64edc676bb98e148987517c332989773fac1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 3 Jun 2026 21:33:46 +0000 Subject: [PATCH 07/40] fix iterator real-input validation for log variance --- src/logstatsexp.jl | 8 +++++++- test/logstatsexp.jl | 4 +++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 330fcf80..5223350e 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -103,13 +103,15 @@ function _logvariance_terms(X) return logsubexp(lse2, 2lse - logn), n end function _logvariance_terms(X, logmean) - R, n = _logsumexp_count((2 * logsubexp(x, logmean) for x in X)) + logmean = _require_real(logmean) + R, n = _logsumexp_count((2 * logsubexp(_require_real(x), logmean) for x in X)) return R, n end function _logsumexp2_count(X) state = iterate(X) state === nothing && throw(ArgumentError("reducing over an empty collection is not allowed")) x, iter_state = state + x = _require_real(x) acc = x acc2 = 2x n = 1 @@ -117,6 +119,7 @@ function _logsumexp2_count(X) state = iterate(X, iter_state) state === nothing && break x, iter_state = state + x = _require_real(x) n += 1 acc = _logsumexp_onepass_op(acc, x) acc2 = _logsumexp_onepass_op(acc2, 2x) @@ -137,3 +140,6 @@ end _reduced_count(X::AbstractArray, R::Number) = length(X) _reduced_count(X::AbstractArray, R::AbstractArray) = length(X) ÷ length(R) + +_require_real(x::Real) = x +_require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index fa3531b8..f8e9a261 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,4 +1,4 @@ -using Test: @test, @testset, @inferred +using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var using LogExpFunctions: logmeanexp, logstdexp, logvarexp @@ -43,6 +43,8 @@ end @test logstdexp(Iterators.Stateful(x)) ≈ log(std(xe)) @test isnan(logvarexp((0.0,))) @test isnan(logstdexp((0.0,))) + @test_throws ArgumentError logvarexp((1.0 + 0.0im, 2.0 + 0.0im)) + @test_throws ArgumentError logstdexp((1.0 + 0.0im, 2.0 + 0.0im)) end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin From 7a886c185c3cea61e4c0c33dbc52d6d30a343376 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 4 Jun 2026 11:57:06 +0000 Subject: [PATCH 08/40] refactor iterator counting and simplify normalization logic --- src/logstatsexp.jl | 58 +++++++++++++++++++++++++++++++++------------- 1 file changed, 42 insertions(+), 16 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 5223350e..4e090015 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -7,7 +7,13 @@ Compute `log(mean(exp, X))`. The result is computed in a numerically stable way. """ function logmeanexp(X) - lse, n = _logsumexp_count(X) + n = _iterator_length(X) + if isnothing(n) + lse, n = _logsumexp_count(X) + else + n == 0 && throw(ArgumentError("reducing over an empty collection is not allowed")) + lse = logsumexp(X) + end return lse - log(_convert_count(lse, n)) end @@ -21,7 +27,8 @@ The result is computed in a numerically stable way. function logmeanexp(X::AbstractArray{<:Number}; dims=:) R = logsumexp(X; dims=dims) n = _reduced_count(X, R) - return _subtract_log_count(R, n) + logn = log(_array_count_type(R, n)) + return R .- logn end """ @@ -33,9 +40,17 @@ Compute `log(var(exp, X; corrected=corrected))`. The result is computed in a numerically stable way. """ function logvarexp(X; corrected::Bool=true, logmean=nothing) - R, n = isnothing(logmean) ? _logvariance_terms(X) : _logvariance_terms(X, logmean) + n = _iterator_length(X) + if isnothing(n) + R, n = isnothing(logmean) ? _logvariance_terms(X) : _logvariance_terms(X, logmean) + else + n == 0 && throw(ArgumentError("reducing over an empty collection is not allowed")) + R = isnothing(logmean) ? _logvariance_terms_reiterable(X, n) : + _logvariance_terms_reiterable(X, n, logmean) + end denom = corrected ? n - 1 : n - return _subtract_log_count(R, denom) + denom == 0 && return oftype(float(R), NaN) + return R - log(_convert_count(R, denom)) end """ @@ -51,7 +66,9 @@ function logvarexp( R = logsumexp(2 * logsubexp.(X, logmean); dims=dims) n = _reduced_count(X, R) denom = corrected ? n - 1 : n - return _subtract_log_count(R, denom) + denom == 0 && return _nan_like(R) + logdenom = log(_array_count_type(R, denom)) + return R .- logdenom end """ @@ -102,6 +119,16 @@ function _logvariance_terms(X) logn = log(_convert_count(lse, n)) return logsubexp(lse2, 2lse - logn), n end +function _logvariance_terms_reiterable(X, n::Integer) + lse = logsumexp(X) + lse2 = logsumexp((2 * _require_real(x) for x in X)) + logn = log(_convert_count(lse, n)) + return logsubexp(lse2, 2lse - logn) +end +function _logvariance_terms_reiterable(X, n::Integer, logmean) + logmean = _require_real(logmean) + return logsumexp((2 * logsubexp(_require_real(x), logmean) for x in X)) +end function _logvariance_terms(X, logmean) logmean = _require_real(logmean) R, n = _logsumexp_count((2 * logsubexp(_require_real(x), logmean) for x in X)) @@ -127,19 +154,18 @@ function _logsumexp2_count(X) return _logsumexp_onepass_result(acc), _logsumexp_onepass_result(acc2), n end -function _subtract_log_count(R::Number, n::Integer) - n == 0 && return oftype(float(R), NaN) - return R - log(_convert_count(R, n)) -end -function _subtract_log_count(R::AbstractArray{<:Number}, n::Integer) - n == 0 && return fill!(R, convert(eltype(R), NaN)) - logn = log(convert(eltype(R), n)) - R .-= logn - return R -end - _reduced_count(X::AbstractArray, R::Number) = length(X) _reduced_count(X::AbstractArray, R::AbstractArray) = length(X) ÷ length(R) +_array_count_type(R::Number, n::Integer) = convert(typeof(R), n) +_array_count_type(R::AbstractArray{<:Number}, n::Integer) = convert(eltype(R), n) +_nan_like(R::Number) = oftype(float(R), NaN) +_nan_like(R::AbstractArray{<:Number}) = fill!(R, convert(eltype(R), NaN)) + +_iterator_length(X) = _iterator_length(Base.IteratorSize(typeof(X)), X) +_iterator_length(::Base.HasLength, X) = length(X) +_iterator_length(::Base.HasShape, X) = length(X) +_iterator_length(::Base.SizeUnknown, X) = nothing +_iterator_length(::Any, X) = nothing _require_real(x::Real) = x _require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) From e76061432098a9ba24fdc60a51eda0ed50578fc4 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 17:43:35 +0200 Subject: [PATCH 09/40] Simplify logstatsexp; add logmeanexp_and_logvarexp / logmeanexp_and_logstdexp Rework the log-mean/var/std-exp reductions around a single primitive (`logsumexp(X)`, `logsumexp(2X)`, count), addressing review feedback in JuliaStats/LogExpFunctions.jl#120 and #77: - Use a single pass over general iterators (works for one-shot iterators such as `Iterators.Stateful`); take the count from `length` for arrays instead of accumulating it, as suggested by @tpapp. - Drop the explicit empty/`NaN` special-casing: for a single element the arithmetic already yields `NaN` (`-Inf - log(0)`), per @tpapp's note. - Remove the `logmean` keyword and the parallel array/iterator variance helpers; everything now flows through `_logvar`/`_logmoments`. - No promotion of `Float32` inputs; full reductions allocate nothing. Add `logmeanexp_and_logvarexp` and `logmeanexp_and_logstdexp` (the `mean_and_var` / `mean_and_std` analogues, also suggested by @tpapp), computing both statistics in a single pass. Tests: extend coverage with the new functions, allocation checks (zero allocations for full reductions over arrays/iterators), and type-stability/`@inferred` checks across `Float32`/`Float64`, `dims`, and `corrected`. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/index.md | 2 + src/LogExpFunctions.jl | 2 +- src/logstatsexp.jl | 237 +++++++++++++++++++++++------------------ test/logstatsexp.jl | 73 ++++++++++++- 4 files changed, 210 insertions(+), 104 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 903acff5..d2d33fbe 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -41,4 +41,6 @@ logit1mexp logmeanexp logvarexp logstdexp +logmeanexp_and_logvarexp +logmeanexp_and_logstdexp ``` diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index 06099f7a..3d8cac3a 100644 --- a/src/LogExpFunctions.jl +++ b/src/LogExpFunctions.jl @@ -9,7 +9,7 @@ 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, logvarexp, logstdexp, + logmeanexp, logvarexp, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp, loglogistic, logitexp, log1mlogistic, logit1mexp include("basicfuns.jl") diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 4e090015..a3273d32 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -1,171 +1,204 @@ +# Numerically stable `log`-of-statistics-of-`exp` reductions. +# +# Everything here is derived from at most two `logsumexp` accumulators, +# +# lse = log ∑ᵢ exp(xᵢ) (first moment) +# lse2 = log ∑ᵢ exp(2xᵢ) (second moment) +# +# together with the number of elements `n`. For a general iterator both +# accumulators are computed in a single pass (important for one-shot iterators +# such as `Iterators.Stateful`); for an `AbstractArray` we reuse the optimized +# `logsumexp` and take `n` from `length`, which is cheap. + """ $(SIGNATURES) -Compute `log(mean(exp, X))`. +Compute `log(mean(exp, X))` in a numerically stable way. `X` should be an iterator of numbers. -The result is computed in a numerically stable way. """ function logmeanexp(X) - n = _iterator_length(X) + n = _known_length(X) if isnothing(n) - lse, n = _logsumexp_count(X) + # length not known ahead of time: count during the single pass + lse, n = _logsumexp_and_count(X) else - n == 0 && throw(ArgumentError("reducing over an empty collection is not allowed")) + iszero(n) && _throw_empty() lse = logsumexp(X) end - return lse - log(_convert_count(lse, n)) + return lse - _log_count(lse, n) end """ $(SIGNATURES) -Compute `log.(mean(exp.(X); dims=dims))`. - -The result is computed in a numerically stable way. +Compute `log.(mean(exp.(X); dims=dims))` in a numerically stable way. """ function logmeanexp(X::AbstractArray{<:Number}; dims=:) - R = logsumexp(X; dims=dims) - n = _reduced_count(X, R) - logn = log(_array_count_type(R, n)) - return R .- logn + lse = logsumexp(X; dims=dims) + return lse .- _log_count(lse, _reduced_count(X, lse)) end """ $(SIGNATURES) -Compute `log(var(exp, X; corrected=corrected))`. +Compute `log(var(exp, X; corrected=corrected))` in a numerically stable way. `X` should be an iterator of real numbers. -The result is computed in a numerically stable way. """ -function logvarexp(X; corrected::Bool=true, logmean=nothing) - n = _iterator_length(X) - if isnothing(n) - R, n = isnothing(logmean) ? _logvariance_terms(X) : _logvariance_terms(X, logmean) - else - n == 0 && throw(ArgumentError("reducing over an empty collection is not allowed")) - R = isnothing(logmean) ? _logvariance_terms_reiterable(X, n) : - _logvariance_terms_reiterable(X, n, logmean) - end - denom = corrected ? n - 1 : n - denom == 0 && return oftype(float(R), NaN) - return R - log(_convert_count(R, denom)) +function logvarexp(X; corrected::Bool=true) + lse, lse2, n = _logmoments(X) + return _logvar(lse, lse2, n, corrected) end """ $(SIGNATURES) -Compute `log.(var(exp.(X); dims=dims, corrected=corrected))`. - -The result is computed in a numerically stable way. +Compute `log.(var(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ -function logvarexp( - X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims) -) - R = logsumexp(2 * logsubexp.(X, logmean); dims=dims) - n = _reduced_count(X, R) - denom = corrected ? n - 1 : n - denom == 0 && return _nan_like(R) - logdenom = log(_array_count_type(R, denom)) - return R .- logdenom +function logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + lse, lse2, n = _logmoments(X, dims) + return _logvar(lse, lse2, n, corrected) end """ $(SIGNATURES) -Compute `log(std(exp, X; corrected=corrected))`. +Compute `log(std(exp, X; corrected=corrected))` in a numerically stable way. `X` should be an iterator of real numbers. -The result is computed in a numerically stable way. """ -function logstdexp(X; corrected::Bool=true, logmean=nothing) - return logvarexp(X; corrected=corrected, logmean=logmean) / 2 +logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected=corrected) / 2 + +""" +$(SIGNATURES) + +Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. +""" +logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = + logvarexp(X; dims=dims, corrected=corrected) ./ 2 + +""" +$(SIGNATURES) + +Compute `(log(mean(exp, X)), log(var(exp, X; corrected=corrected)))` in a numerically +stable way, using a single pass over the data. + +`X` should be an iterator of real numbers. +""" +function logmeanexp_and_logvarexp(X; corrected::Bool=true) + lse, lse2, n = _logmoments(X) + return lse - _log_count(lse, n), _logvar(lse, lse2, n, corrected) end """ $(SIGNATURES) -Compute `log.(std(exp.(X); dims=dims, corrected=corrected))`. +Compute `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))` in a numerically +stable way. +""" +function logmeanexp_and_logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + lse, lse2, n = _logmoments(X, dims) + return lse .- _log_count(lse, n), _logvar(lse, lse2, n, corrected) +end + +""" +$(SIGNATURES) + +Compute `(log(mean(exp, X)), log(std(exp, X; corrected=corrected)))` in a numerically +stable way, using a single pass over the data. + +`X` should be an iterator of real numbers. +""" +function logmeanexp_and_logstdexp(X; corrected::Bool=true) + logmean, logvar = logmeanexp_and_logvarexp(X; corrected=corrected) + return logmean, logvar / 2 +end + +""" +$(SIGNATURES) -The result is computed in a numerically stable way. +Compute `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))` in a numerically +stable way. """ -function logstdexp( - X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims=dims) -) - return logvarexp(X; dims=dims, corrected=corrected, logmean=logmean) / 2 +function logmeanexp_and_logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + logmean, logvar = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) + return logmean, logvar ./ 2 +end + +# ---- internal helpers ---- + +_throw_empty() = throw(ArgumentError("reducing over an empty collection is not allowed")) + +# `log(n)` in the (real) floating point type of `R`, so that no unwanted promotion of +# `R` (e.g. `Float32` to `Float64`) takes place. Works for scalar and array `R` alike. +# `n == 0` gives `log(0) == -Inf`, which produces the expected `NaN`/`-Inf` results. +_log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) + +# number of elements reduced into each entry of `R` +_reduced_count(X::AbstractArray, R) = length(X) ÷ length(R) + +_require_real(x::Real) = x +_require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) + +_known_length(X) = _known_length(Base.IteratorSize(typeof(X)), X) +_known_length(::Union{Base.HasLength,Base.HasShape}, X) = length(X) +_known_length(_, X) = nothing + +# `log(var)` from the raw log-moments. The squared-deviation sum is +# `∑ᵢ (exp(xᵢ) - mean)² = ∑ᵢ exp(2xᵢ) - (∑ᵢ exp(xᵢ))² / n`, i.e. +# `logsubexp(lse2, 2 * lse - log(n))` in log space. For a single element +# (`n - corrected == 0`) the numerator is `-Inf` and `log` of the denominator is also +# `-Inf`, so the result is `NaN`, matching `var`. +function _logvar(lse, lse2, n::Integer, corrected::Bool) + logn = _log_count(lse2, n) + logdenom = _log_count(lse2, corrected ? n - 1 : n) + return @. logsubexp(lse2, 2 * lse - logn) - logdenom end -function _logsumexp_count(X) - state = iterate(X) - state === nothing && throw(ArgumentError("reducing over an empty collection is not allowed")) - x, iter_state = state +# Single pass over a general iterator: (logsumexp(X), count). +function _logsumexp_and_count(X) + next = iterate(X) + isnothing(next) && _throw_empty() + x, state = next acc = x n = 1 while true - state = iterate(X, iter_state) - state === nothing && break - x, iter_state = state - n += 1 + next = iterate(X, state) + isnothing(next) && break + x, state = next acc = _logsumexp_onepass_op(acc, x) + n += 1 end return _logsumexp_onepass_result(acc), n end -_convert_count(x, n::Integer) = convert(typeof(x), n) - -function _logvariance_terms(X) - lse, lse2, n = _logsumexp2_count(X) - logn = log(_convert_count(lse, n)) - return logsubexp(lse2, 2lse - logn), n -end -function _logvariance_terms_reiterable(X, n::Integer) - lse = logsumexp(X) - lse2 = logsumexp((2 * _require_real(x) for x in X)) - logn = log(_convert_count(lse, n)) - return logsubexp(lse2, 2lse - logn) -end -function _logvariance_terms_reiterable(X, n::Integer, logmean) - logmean = _require_real(logmean) - return logsumexp((2 * logsubexp(_require_real(x), logmean) for x in X)) -end -function _logvariance_terms(X, logmean) - logmean = _require_real(logmean) - R, n = _logsumexp_count((2 * logsubexp(_require_real(x), logmean) for x in X)) - return R, n -end -function _logsumexp2_count(X) - state = iterate(X) - state === nothing && throw(ArgumentError("reducing over an empty collection is not allowed")) - x, iter_state = state +# Single pass over a general iterator of reals: (logsumexp(X), logsumexp(2X), count). +function _logmoments(X) + next = iterate(X) + isnothing(next) && _throw_empty() + x, state = next x = _require_real(x) acc = x acc2 = 2x n = 1 while true - state = iterate(X, iter_state) - state === nothing && break - x, iter_state = state + next = iterate(X, state) + isnothing(next) && break + x, state = next x = _require_real(x) - n += 1 acc = _logsumexp_onepass_op(acc, x) acc2 = _logsumexp_onepass_op(acc2, 2x) + n += 1 end return _logsumexp_onepass_result(acc), _logsumexp_onepass_result(acc2), n end -_reduced_count(X::AbstractArray, R::Number) = length(X) -_reduced_count(X::AbstractArray, R::AbstractArray) = length(X) ÷ length(R) -_array_count_type(R::Number, n::Integer) = convert(typeof(R), n) -_array_count_type(R::AbstractArray{<:Number}, n::Integer) = convert(eltype(R), n) -_nan_like(R::Number) = oftype(float(R), NaN) -_nan_like(R::AbstractArray{<:Number}) = fill!(R, convert(eltype(R), NaN)) - -_iterator_length(X) = _iterator_length(Base.IteratorSize(typeof(X)), X) -_iterator_length(::Base.HasLength, X) = length(X) -_iterator_length(::Base.HasShape, X) = length(X) -_iterator_length(::Base.SizeUnknown, X) = nothing -_iterator_length(::Any, X) = nothing - -_require_real(x::Real) = x -_require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) +# Array version: reuse the optimized `logsumexp` and take the count from `length`. +function _logmoments(X::AbstractArray{<:Real}, dims) + # full reduction: a single pass avoids allocating the `2 .* X` temporary + dims === Colon() && return _logmoments(X) + lse = logsumexp(X; dims=dims) + lse2 = logsumexp(2 .* X; dims=dims) + return lse, lse2, _reduced_count(X, lse) +end diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index f8e9a261..a51f2545 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,6 +1,12 @@ using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var -using LogExpFunctions: logmeanexp, logstdexp, logvarexp +using LogExpFunctions: logmeanexp, logstdexp, logvarexp, + logmeanexp_and_logvarexp, logmeanexp_and_logstdexp + +# Count heap allocations of `f(x)` after warming it up. `f` and `x` are passed as +# arguments so that they are concretely typed inside this function (avoiding spurious +# allocations from captured, boxed locals). +allocations(f, x) = (f(x); @allocated f(x)) @testset "logmeanexp, logvarexp, logstdexp arrays" begin for T in (Float32, Float64) @@ -45,6 +51,8 @@ end @test isnan(logstdexp((0.0,))) @test_throws ArgumentError logvarexp((1.0 + 0.0im, 2.0 + 0.0im)) @test_throws ArgumentError logstdexp((1.0 + 0.0im, 2.0 + 0.0im)) + @test_throws ArgumentError logmeanexp(()) + @test_throws ArgumentError logvarexp(()) end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin @@ -67,3 +75,66 @@ end @test isnan(logvarexp(Xsingleton; dims=:, corrected=true)) @test isnan(logstdexp(Xsingleton; dims=:, corrected=true)) end + +@testset "logmeanexp_and_logvarexp, logmeanexp_and_logstdexp" begin + for T in (Float32, Float64) + X = randn(T, 5, 3, 2) + for dims in (2, (1, 2), :), corrected in (true, false) + m, v = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) + @test m ≈ logmeanexp(X; dims=dims) + @test v ≈ logvarexp(X; dims=dims, corrected=corrected) + m2, s = logmeanexp_and_logstdexp(X; dims=dims, corrected=corrected) + @test m2 ≈ logmeanexp(X; dims=dims) + @test s ≈ logstdexp(X; dims=dims, corrected=corrected) + end + # results match the reference statistics directly + @test all(logmeanexp_and_logvarexp(X) .≈ (log(mean(exp, X)), log(var(exp.(X))))) + @test all(logmeanexp_and_logstdexp(X) .≈ (log(mean(exp, X)), log(std(exp.(X))))) + end + + # iterators (single pass, including one-shot iterators) + x = randn(Float32, 20) + xt = Tuple(x) + xe = exp.(x) + @test all(@inferred(logmeanexp_and_logvarexp(xt)) .≈ (log(mean(exp, xt)), log(var(xe)))) + @test all(@inferred(logmeanexp_and_logstdexp(xt)) .≈ (log(mean(exp, xt)), log(std(xe)))) + @test all(logmeanexp_and_logvarexp(Iterators.Stateful(x)) .≈ (log(mean(exp, x)), log(var(xe)))) + @test all(logmeanexp_and_logstdexp(Iterators.Stateful(x)) .≈ (log(mean(exp, x)), log(std(xe)))) + + # edge cases + @test isnan(last(logmeanexp_and_logvarexp((0.0,)))) + @test isnan(last(logmeanexp_and_logstdexp((0.0,)))) + @test_throws ArgumentError logmeanexp_and_logvarexp((1.0 + 0.0im,)) + @test_throws ArgumentError logmeanexp_and_logstdexp((1.0 + 0.0im,)) +end + +@testset "type stability and inference" begin + X = randn(Float32, 5, 3, 2) + xt = Tuple(randn(Float32, 20)) + for dims in (1, (2, 3), :) + @test @inferred(logmeanexp_and_logvarexp(X; dims=dims)) isa Tuple + @test @inferred(logmeanexp_and_logstdexp(X; dims=dims)) isa Tuple + end + @test @inferred(logmeanexp_and_logvarexp(xt)) isa NTuple{2,Float32} + @test @inferred(logmeanexp_and_logstdexp(xt)) isa NTuple{2,Float32} + + # no Float64 promotion for Float32 inputs + m, v = logmeanexp_and_logvarexp(X; dims=2) + @test eltype(m) == Float32 && eltype(v) == Float32 + @test typeof(@inferred(logmeanexp_and_logvarexp(X))) == Tuple{Float32,Float32} +end + +@testset "allocations" begin + for T in (Float32, Float64) + v = randn(T, 1000) + tup = Tuple(randn(T, 20)) + # full reductions over arrays / iterators allocate nothing + @test allocations(logmeanexp, v) == 0 + @test allocations(logvarexp, v) == 0 + @test allocations(logstdexp, v) == 0 + @test allocations(logmeanexp_and_logvarexp, v) == 0 + @test allocations(logmeanexp_and_logstdexp, v) == 0 + @test allocations(logvarexp, tup) == 0 + @test allocations(logmeanexp_and_logvarexp, tup) == 0 + end +end From 3842a7e9ce04abe491cb932713ce234d655dad3c Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 18:31:21 +0200 Subject: [PATCH 10/40] Use the numerically stable centered variance; fix min-Julia CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Numerical robustness fix. The previous single-pass raw-moment variance `logsubexp(∑exp(2xᵢ), (∑exp(xᵢ))²/n)` cancels catastrophically when the variance is small relative to the mean — verified against BigFloat references it loses all accuracy for tight clusters and even returns `Inf` for nearly-equal inputs. Switch to the centered formula log var = logsumexp(2 * logsubexp(xᵢ, logmean)) - log(n - corrected) which is accurate across the board (tight clusters, overflow at exp(±700), huge dynamic range). The mean is reused to center the variance, so the combined `logmeanexp_and_logvarexp` is still cheaper than separate calls. Because the variance now needs two passes (mean, then deviations), general iterators are materialized once with `collect`; re-iterable containers (arrays, tuples, ranges) are traversed in place. `IteratorSize` is not used to detect re-iterability — `Iterators.Stateful` reports `HasLength` on Julia 1.10 but is single-use, which is what broke the previous version. Tests: - add a `numerical robustness` testset comparing against BigFloat on the hard cases above (would fail with the raw-moment formula). - the allocation tests now assert allocations do not scale with input size (a small constant), instead of exactly zero: Julia 1.10's optimizer keeps a tiny constant allocation that 1.12 elides, which was the CI failure on `Julia min-patch`. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 131 +++++++++++++++++++++----------------------- test/logstatsexp.jl | 53 ++++++++++++++---- 2 files changed, 105 insertions(+), 79 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index a3273d32..ceae2fd6 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -1,14 +1,14 @@ # Numerically stable `log`-of-statistics-of-`exp` reductions. # -# Everything here is derived from at most two `logsumexp` accumulators, +# The mean is `logsumexp(X) - log(n)`. The variance uses the *centered* formula # -# lse = log ∑ᵢ exp(xᵢ) (first moment) -# lse2 = log ∑ᵢ exp(2xᵢ) (second moment) +# log var = logsumexp(2 * logsubexp(xᵢ, logmean)) - log(n - corrected) # -# together with the number of elements `n`. For a general iterator both -# accumulators are computed in a single pass (important for one-shot iterators -# such as `Iterators.Stateful`); for an `AbstractArray` we reuse the optimized -# `logsumexp` and take `n` from `length`, which is cheap. +# i.e. the log of the sum of squared deviations `∑ᵢ (exp(xᵢ) - mean)²`, divided by the +# count. Centering is essential for numerical stability: the raw-moment alternative +# `logsubexp(∑exp(2xᵢ), (∑exp(xᵢ))²/n)` cancels catastrophically when the variance is +# small relative to the mean (and can even overflow to `Inf` for nearly-equal inputs). +# Computing the variance therefore needs the mean first, hence the two helpers below. """ $(SIGNATURES) @@ -42,71 +42,75 @@ end """ $(SIGNATURES) -Compute `log(var(exp, X; corrected=corrected))` in a numerically stable way. +Compute `(log(mean(exp, X)), log(var(exp, X; corrected=corrected)))` in a numerically +stable way. Computing the two together is cheaper than calling [`logmeanexp`](@ref) and +[`logvarexp`](@ref) separately, since the mean is reused to center the variance. `X` should be an iterator of real numbers. """ -function logvarexp(X; corrected::Bool=true) - lse, lse2, n = _logmoments(X) - return _logvar(lse, lse2, n, corrected) +function logmeanexp_and_logvarexp(X; corrected::Bool=true) + xs = _materialize(X) + logmean = logmeanexp(xs) + logsqdev = _centered_logsqdev(xs, logmean) + return logmean, _finish_logvar(logsqdev, length(xs), corrected) end """ $(SIGNATURES) -Compute `log.(var(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. +Compute `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))` in a +numerically stable way, reusing the mean to center the variance. """ -function logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) - lse, lse2, n = _logmoments(X, dims) - return _logvar(lse, lse2, n, corrected) +function logmeanexp_and_logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + logmean = logmeanexp(X; dims=dims) + return logmean, _centered_logvar(X, logmean, dims, corrected) end +# dispatch on `dims` so the return type is concrete (no `Union` of scalar/array results) +_centered_logvar(X, logmean, ::Colon, corrected::Bool) = # scalar reduction, no temporary + _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) +_centered_logvar(X, logmean, dims, corrected::Bool) = + _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) + """ $(SIGNATURES) -Compute `log(std(exp, X; corrected=corrected))` in a numerically stable way. +Compute `log(var(exp, X; corrected=corrected))` in a numerically stable way. `X` should be an iterator of real numbers. """ -logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected=corrected) / 2 +logvarexp(X; corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; corrected=corrected)) """ $(SIGNATURES) -Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. +Compute `log.(var(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ -logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = - logvarexp(X; dims=dims, corrected=corrected) ./ 2 +logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = + last(logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected)) """ $(SIGNATURES) -Compute `(log(mean(exp, X)), log(var(exp, X; corrected=corrected)))` in a numerically -stable way, using a single pass over the data. +Compute `log(std(exp, X; corrected=corrected))` in a numerically stable way. `X` should be an iterator of real numbers. """ -function logmeanexp_and_logvarexp(X; corrected::Bool=true) - lse, lse2, n = _logmoments(X) - return lse - _log_count(lse, n), _logvar(lse, lse2, n, corrected) -end +logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected=corrected) / 2 """ $(SIGNATURES) -Compute `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))` in a numerically -stable way. +Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ -function logmeanexp_and_logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) - lse, lse2, n = _logmoments(X, dims) - return lse .- _log_count(lse, n), _logvar(lse, lse2, n, corrected) -end +logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = + logvarexp(X; dims=dims, corrected=corrected) ./ 2 """ $(SIGNATURES) Compute `(log(mean(exp, X)), log(std(exp, X; corrected=corrected)))` in a numerically -stable way, using a single pass over the data. +stable way, reusing the mean to center the variance. `X` should be an iterator of real numbers. """ @@ -118,8 +122,8 @@ end """ $(SIGNATURES) -Compute `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))` in a numerically -stable way. +Compute `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))` in a +numerically stable way, reusing the mean to center the variance. """ function logmeanexp_and_logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) logmean, logvar = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) @@ -145,60 +149,49 @@ _known_length(X) = _known_length(Base.IteratorSize(typeof(X)), X) _known_length(::Union{Base.HasLength,Base.HasShape}, X) = length(X) _known_length(_, X) = nothing -# `log(var)` from the raw log-moments. The squared-deviation sum is -# `∑ᵢ (exp(xᵢ) - mean)² = ∑ᵢ exp(2xᵢ) - (∑ᵢ exp(xᵢ))² / n`, i.e. -# `logsubexp(lse2, 2 * lse - log(n))` in log space. For a single element -# (`n - corrected == 0`) the numerator is `-Inf` and `log` of the denominator is also -# `-Inf`, so the result is `NaN`, matching `var`. -function _logvar(lse, lse2, n::Integer, corrected::Bool) - logn = _log_count(lse2, n) - logdenom = _log_count(lse2, corrected ? n - 1 : n) - return @. logsubexp(lse2, 2 * lse - logn) - logdenom -end - -# Single pass over a general iterator: (logsumexp(X), count). -function _logsumexp_and_count(X) +# Variance is centered, so we need to traverse the data twice (mean, then deviations). +# Known re-iterable containers are traversed in place; any other iterator is materialized +# once, so that one-shot iterators (e.g. `Iterators.Stateful`) are handled correctly. +# (`IteratorSize` cannot be used to detect re-iterability: `Stateful` reports `HasLength` +# on some Julia versions yet is single-use.) +_materialize(X) = collect(X) +_materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X + +# log of the sum of squared deviations, `logsumexp(2 * logsubexp(xᵢ, logmean))`, +# accumulated in a single pass without allocating an intermediate array. +function _centered_logsqdev(X, logmean) next = iterate(X) isnothing(next) && _throw_empty() x, state = next - acc = x - n = 1 + acc = 2 * logsubexp(_require_real(x), logmean) while true next = iterate(X, state) isnothing(next) && break x, state = next - acc = _logsumexp_onepass_op(acc, x) - n += 1 + acc = _logsumexp_onepass_op(acc, 2 * logsubexp(_require_real(x), logmean)) end - return _logsumexp_onepass_result(acc), n + return _logsumexp_onepass_result(acc) end -# Single pass over a general iterator of reals: (logsumexp(X), logsumexp(2X), count). -function _logmoments(X) +# divide the squared-deviation sum by the count, in log space. For a single element +# (`n - corrected == 0`) the numerator is `-Inf` and `log` of the denominator is also +# `-Inf`, so the result is `NaN`, matching `var`. +_finish_logvar(logsqdev, n::Integer, corrected::Bool) = + logsqdev .- _log_count(logsqdev, corrected ? n - 1 : n) + +# Single pass over a general iterator: (logsumexp(X), count). +function _logsumexp_and_count(X) next = iterate(X) isnothing(next) && _throw_empty() x, state = next - x = _require_real(x) acc = x - acc2 = 2x n = 1 while true next = iterate(X, state) isnothing(next) && break x, state = next - x = _require_real(x) acc = _logsumexp_onepass_op(acc, x) - acc2 = _logsumexp_onepass_op(acc2, 2x) n += 1 end - return _logsumexp_onepass_result(acc), _logsumexp_onepass_result(acc2), n -end - -# Array version: reuse the optimized `logsumexp` and take the count from `length`. -function _logmoments(X::AbstractArray{<:Real}, dims) - # full reduction: a single pass avoids allocating the `2 .* X` temporary - dims === Colon() && return _logmoments(X) - lse = logsumexp(X; dims=dims) - lse2 = logsumexp(2 .* X; dims=dims) - return lse, lse2, _reduced_count(X, lse) + return _logsumexp_onepass_result(acc), n end diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index a51f2545..0f93accb 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -125,16 +125,49 @@ end end @testset "allocations" begin + # Full reductions allocate at most a small constant (no per-element / O(n) temporary): + # the allocation count must not grow with the input size. for T in (Float32, Float64) - v = randn(T, 1000) - tup = Tuple(randn(T, 20)) - # full reductions over arrays / iterators allocate nothing - @test allocations(logmeanexp, v) == 0 - @test allocations(logvarexp, v) == 0 - @test allocations(logstdexp, v) == 0 - @test allocations(logmeanexp_and_logvarexp, v) == 0 - @test allocations(logmeanexp_and_logstdexp, v) == 0 - @test allocations(logvarexp, tup) == 0 - @test allocations(logmeanexp_and_logvarexp, tup) == 0 + for f in (logmeanexp, logvarexp, logstdexp, + logmeanexp_and_logvarexp, logmeanexp_and_logstdexp) + @test allocations(f, randn(T, 10_000)) == allocations(f, randn(T, 100)) + end + # genuinely allocation-free paths + @test allocations(logmeanexp, randn(T, 10_000)) == 0 + @test allocations(logvarexp, Tuple(randn(T, 20))) == 0 + end +end + +@testset "numerical robustness" begin + # Compare against high-precision (BigFloat) references on hard cases: tight clusters + # (var ≪ mean², where a raw second-moment formula cancels catastrophically) and + # values large/small enough that exp would over-/under-flow Float64. + setprecision(BigFloat, 256) do + refmean(x) = Float64(log(mean(exp.(big.(x))))) + refvar(x; corrected=true) = Float64(log(var(exp.(big.(x)); corrected=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 (e.g. logmeanexp([0, 1e-6]) ≈ 5e-7), where + # an inherent cancellation makes the relative error meaningless + @test logmeanexp(x) ≈ refmean(x) rtol = 1e-9 atol = 1e-10 + for corrected in (true, false) + @test logvarexp(x; corrected=corrected) ≈ refvar(x; corrected=corrected) rtol = 1e-8 atol = 1e-9 + @test logstdexp(x; corrected=corrected) ≈ refvar(x; corrected=corrected) / 2 rtol = 1e-8 atol = 1e-9 + end + end + # array and one-shot-iterator paths agree even in the cancellation regime + x = 1.0 .+ 1e-5 .* randn(100) + @test logvarexp(x) ≈ logvarexp(Iterators.Stateful(x)) rtol = 1e-10 + m, v = logmeanexp_and_logvarexp(Iterators.Stateful(x)) + @test m ≈ logmeanexp(x) rtol = 1e-10 + @test v ≈ logvarexp(x) rtol = 1e-10 end end From 6789ee65bcc88c7db8f0e1a98f10e9e0937b6e84 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 20:21:45 +0200 Subject: [PATCH 11/40] Use scalar `/ 2` instead of broadcast `./ 2` for std `AbstractArray / Number` is defined in Base, so the dot is redundant when dividing by a scalar. This makes the array `logstdexp` / `logmeanexp_and_logstdexp` methods consistent with their scalar siblings. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index ceae2fd6..7f61b7c4 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -104,7 +104,7 @@ $(SIGNATURES) Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = - logvarexp(X; dims=dims, corrected=corrected) ./ 2 + logvarexp(X; dims=dims, corrected=corrected) / 2 """ $(SIGNATURES) @@ -127,7 +127,7 @@ numerically stable way, reusing the mean to center the variance. """ function logmeanexp_and_logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) logmean, logvar = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) - return logmean, logvar ./ 2 + return logmean, logvar / 2 end # ---- internal helpers ---- From 8c421b5ff058dabe3cb7d1ca963c01c43ebea76b Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 21:41:44 +0200 Subject: [PATCH 12/40] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/logstatsexp.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 7f61b7c4..649d6ea5 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -35,6 +35,7 @@ $(SIGNATURES) Compute `log.(mean(exp.(X); dims=dims))` in a numerically stable way. """ function logmeanexp(X::AbstractArray{<:Number}; dims=:) + dims isa Colon && isempty(X) && _throw_empty() lse = logsumexp(X; dims=dims) return lse .- _log_count(lse, _reduced_count(X, lse)) end From 73f2e554081cb1652e1a17805727a6928e5d3621 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 21:42:00 +0200 Subject: [PATCH 13/40] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/logstatsexp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 0f93accb..4668ddc1 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -52,8 +52,8 @@ end @test_throws ArgumentError logvarexp((1.0 + 0.0im, 2.0 + 0.0im)) @test_throws ArgumentError logstdexp((1.0 + 0.0im, 2.0 + 0.0im)) @test_throws ArgumentError logmeanexp(()) + @test_throws ArgumentError logmeanexp(Float64[]) @test_throws ArgumentError logvarexp(()) -end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin X = randn(Float32, 5, 3, 2) From 9fa4627021dd919b1701f7983ba318a641a52049 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 22:12:25 +0200 Subject: [PATCH 14/40] Fuse dims variance into logsumexp; repair empty-array testset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback on the `dims != :` variance path: it no longer materializes a full-size `2 .* logsubexp.(X, logmean)` temporary before reducing. Instead `_LazyLogSqDev` is a read-only `AbstractArray` view of that broadcast, so `logsumexp(...; dims)` reduces it on the fly and only allocates the (smaller) reduced output. A bare `Broadcasted` can't be reduced along `dims` (no `axes(_, i)`), hence the thin wrapper. The elementwise term `2 * logsubexp(xᵢ, logmean)` is factored into `_logsqdev_term`, shared by the scalar single-pass loop and the lazy view. Also: - repair the "iterators" `@testset` whose closing `end` was dropped when the empty-array regression test was added (the file no longer parsed); - add an empty-array `logvarexp(Float64[])` regression test; - add allocation regression tests asserting the `dims` reductions allocate only the fixed-size reduced output, not an O(length(X)) temporary. Verified on Julia 1.10.11 and 1.12.6: full Pkg.test (incl. JET + Aqua) passes. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 28 +++++++++++++++++++++++++--- test/logstatsexp.jl | 12 ++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 649d6ea5..8ec2e1c3 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -71,7 +71,7 @@ end _centered_logvar(X, logmean, ::Colon, corrected::Bool) = # scalar reduction, no temporary _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) _centered_logvar(X, logmean, dims, corrected::Bool) = - _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) + _finish_logvar(logsumexp(_logsqdev_lazy(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) """ $(SIGNATURES) @@ -158,22 +158,44 @@ _known_length(_, X) = nothing _materialize(X) = collect(X) _materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X +# log of the (centered) squared deviation of a single point, `2 * logsubexp(xᵢ, logmean)` +# (i.e. `log((exp(xᵢ) - mean)^2)`). This is the term summed to form the variance. +_logsqdev_term(x, logmean) = 2 * logsubexp(_require_real(x), logmean) + # log of the sum of squared deviations, `logsumexp(2 * logsubexp(xᵢ, logmean))`, # accumulated in a single pass without allocating an intermediate array. function _centered_logsqdev(X, logmean) next = iterate(X) isnothing(next) && _throw_empty() x, state = next - acc = 2 * logsubexp(_require_real(x), logmean) + acc = _logsqdev_term(x, logmean) while true next = iterate(X, state) isnothing(next) && break x, state = next - acc = _logsumexp_onepass_op(acc, 2 * logsubexp(_require_real(x), logmean)) + acc = _logsumexp_onepass_op(acc, _logsqdev_term(x, logmean)) end return _logsumexp_onepass_result(acc) end +# For a `dims`-reduction we cannot accumulate in a single scalar pass, but we still avoid +# materializing the full-size `2 .* logsubexp.(X, logmean)`: `_LazyLogSqDev` is a read-only +# `AbstractArray` view of that broadcast, so `logsumexp(...; dims)` reduces it on the fly +# and only allocates the (smaller) reduced result. A bare `Broadcasted` cannot be used +# directly, as it does not implement the `axes(_, i)` that `dims` reductions require. +struct _LazyLogSqDev{T,N,B<:Base.Broadcast.Broadcasted} <: AbstractArray{T,N} + bc::B +end +function _logsqdev_lazy(X, logmean) + bc = Base.Broadcast.instantiate(Base.broadcasted(_logsqdev_term, X, logmean)) + T = Base.Broadcast.combine_eltypes(bc.f, bc.args) + return _LazyLogSqDev{T,ndims(bc),typeof(bc)}(bc) +end +Base.size(A::_LazyLogSqDev) = map(length, axes(A.bc)) +Base.axes(A::_LazyLogSqDev) = axes(A.bc) +Base.IndexStyle(::Type{<:_LazyLogSqDev}) = IndexCartesian() +Base.@propagate_inbounds Base.getindex(A::_LazyLogSqDev, I::Int...) = A.bc[I...] + # divide the squared-deviation sum by the count, in log space. For a single element # (`n - corrected == 0`) the numerator is `-Inf` and `log` of the denominator is also # `-Inf`, so the result is `NaN`, matching `var`. diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 4668ddc1..05c04c84 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -7,6 +7,9 @@ using LogExpFunctions: logmeanexp, logstdexp, logvarexp, # arguments so that they are concretely typed inside this function (avoiding spurious # allocations from captured, boxed locals). allocations(f, x) = (f(x); @allocated f(x)) +# Same, for the `dims` keyword form (`x` is preallocated by the caller, so only the +# allocations of `f(x; dims=1)` itself are counted). +allocations_dims(f, x) = (f(x; dims=1); @allocated f(x; dims=1)) @testset "logmeanexp, logvarexp, logstdexp arrays" begin for T in (Float32, Float64) @@ -54,6 +57,8 @@ end @test_throws ArgumentError logmeanexp(()) @test_throws ArgumentError logmeanexp(Float64[]) @test_throws ArgumentError logvarexp(()) + @test_throws ArgumentError logvarexp(Float64[]) +end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin X = randn(Float32, 5, 3, 2) @@ -135,6 +140,13 @@ end # genuinely allocation-free paths @test allocations(logmeanexp, randn(T, 10_000)) == 0 @test allocations(logvarexp, Tuple(randn(T, 20))) == 0 + # `dims` reductions allocate only the (fixed-size) reduced output, never a + # temporary the size of the input: the variance is fused into `logsumexp` lazily. + # Reducing along `dims=1` gives the same `(1, 4)` output for both sizes, so a + # constant allocation count confirms no O(length(X)) intermediate is built. + for f in (logvarexp, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp) + @test allocations_dims(f, randn(T, 10_000, 4)) == allocations_dims(f, randn(T, 100, 4)) + end end end From dadf0d92d1c259d4d2623e2a9afc5546601506a8 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 23:04:04 +0200 Subject: [PATCH 15/40] Fix dims-reduction correctness bugs found in review; drop lazy wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A self-review of this branch surfaced several correctness bugs, all on the `dims` array paths. Fixes: - Revert the `_LazyLogSqDev` lazy-reduction wrapper back to the plain `logsumexp(2 .* logsubexp.(X, logmean); dims)`. The wrapper forwarded raw Cartesian indices to the wrapped `Broadcasted`, which silently returned WRONG variance/std for arrays with non-1-based axes (e.g. OffsetArrays — mean was correct, so mean and variance became mutually inconsistent), and threw a MethodError for abstract element types (`combine_eltypes` → `Any`). The factor of 2 cannot be pulled out of `logsumexp` (`log Σ exp(2y) ≠ 2 log Σ exp(y)`), and the materialized broadcast is correct and simpler; one temporary the size of the input is the same thing `Statistics.var(exp.(X); dims)` allocates. - Empty reduction with `corrected=true` no longer throws `DomainError` from `log(-1)`; `_finish_logvar` clamps the (corrected) count to ≥ 0 so an empty reduction yields `NaN`, matching `Statistics.var` and `logmeanexp`. - Empty along a non-reduced dimension no longer throws `DivideError` (0 ÷ 0); `_reduced_count` returns 0 when the reduced result is empty, giving an empty array of the reduced shape. - Complex arrays with `dims` now throw the clear "require real inputs" ArgumentError (the variance/std array methods accept `<:Number` and reject non-real eltype up front via `_require_real_array`) instead of a confusing MethodError; behavior now matches the no-`dims` path. Adds an "edge-case regressions" testset covering all of the above (OffsetArrays, abstract eltype, empty reduced/non-reduced dims, complex + dims). Removes the dims allocation test that asserted the (now reverted) lazy fusion. Verified on Julia 1.10.11 and 1.12.6. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 50 +++++++++++++++--------------------- test/logstatsexp.jl | 62 +++++++++++++++++++++++++++++++++++++-------- 2 files changed, 73 insertions(+), 39 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 8ec2e1c3..c487c5b7 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -62,7 +62,8 @@ $(SIGNATURES) Compute `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))` in a numerically stable way, reusing the mean to center the variance. """ -function logmeanexp_and_logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) +function logmeanexp_and_logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) + _require_real_array(X) logmean = logmeanexp(X; dims=dims) return logmean, _centered_logvar(X, logmean, dims, corrected) end @@ -71,7 +72,7 @@ end _centered_logvar(X, logmean, ::Colon, corrected::Bool) = # scalar reduction, no temporary _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) _centered_logvar(X, logmean, dims, corrected::Bool) = - _finish_logvar(logsumexp(_logsqdev_lazy(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) + _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) """ $(SIGNATURES) @@ -87,7 +88,7 @@ $(SIGNATURES) Compute `log.(var(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ -logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = +logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected)) """ @@ -104,7 +105,7 @@ $(SIGNATURES) Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. """ -logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = +logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = logvarexp(X; dims=dims, corrected=corrected) / 2 """ @@ -126,7 +127,7 @@ $(SIGNATURES) Compute `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))` in a numerically stable way, reusing the mean to center the variance. """ -function logmeanexp_and_logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) +function logmeanexp_and_logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) logmean, logvar = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) return logmean, logvar / 2 end @@ -140,12 +141,20 @@ _throw_empty() = throw(ArgumentError("reducing over an empty collection is not a # `n == 0` gives `log(0) == -Inf`, which produces the expected `NaN`/`-Inf` results. _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) -# number of elements reduced into each entry of `R` -_reduced_count(X::AbstractArray, R) = length(X) ÷ length(R) +# number of elements reduced into each entry of `R`. When `R` is empty (the reduction +# produced no cells, e.g. `X` is empty along a dimension that is not being reduced) the +# count is irrelevant — the result is empty — so avoid dividing by zero. +_reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) _require_real(x::Real) = x _require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) +# variance/std require real inputs; reject a non-real element type up front (with the same +# message as `_require_real`) so the array `dims` paths fail cleanly instead of hitting a +# `MethodError` deep inside `logsubexp`. +_require_real_array(X::AbstractArray{<:Real}) = X +_require_real_array(X::AbstractArray) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) + _known_length(X) = _known_length(Base.IteratorSize(typeof(X)), X) _known_length(::Union{Base.HasLength,Base.HasShape}, X) = length(X) _known_length(_, X) = nothing @@ -178,29 +187,12 @@ function _centered_logsqdev(X, logmean) return _logsumexp_onepass_result(acc) end -# For a `dims`-reduction we cannot accumulate in a single scalar pass, but we still avoid -# materializing the full-size `2 .* logsubexp.(X, logmean)`: `_LazyLogSqDev` is a read-only -# `AbstractArray` view of that broadcast, so `logsumexp(...; dims)` reduces it on the fly -# and only allocates the (smaller) reduced result. A bare `Broadcasted` cannot be used -# directly, as it does not implement the `axes(_, i)` that `dims` reductions require. -struct _LazyLogSqDev{T,N,B<:Base.Broadcast.Broadcasted} <: AbstractArray{T,N} - bc::B -end -function _logsqdev_lazy(X, logmean) - bc = Base.Broadcast.instantiate(Base.broadcasted(_logsqdev_term, X, logmean)) - T = Base.Broadcast.combine_eltypes(bc.f, bc.args) - return _LazyLogSqDev{T,ndims(bc),typeof(bc)}(bc) -end -Base.size(A::_LazyLogSqDev) = map(length, axes(A.bc)) -Base.axes(A::_LazyLogSqDev) = axes(A.bc) -Base.IndexStyle(::Type{<:_LazyLogSqDev}) = IndexCartesian() -Base.@propagate_inbounds Base.getindex(A::_LazyLogSqDev, I::Int...) = A.bc[I...] - -# divide the squared-deviation sum by the count, in log space. For a single element -# (`n - corrected == 0`) the numerator is `-Inf` and `log` of the denominator is also -# `-Inf`, so the result is `NaN`, matching `var`. +# divide the squared-deviation sum by the count, in log space. When the (corrected) count +# is non-positive — a single element with `corrected=true`, or an empty reduction — the +# numerator and `log` of the (clamped-to-zero) denominator are both `-Inf`, giving `NaN`, +# matching `var`. Clamping to zero also avoids `log(-1)` for an empty `corrected=true` case. _finish_logvar(logsqdev, n::Integer, corrected::Bool) = - logsqdev .- _log_count(logsqdev, corrected ? n - 1 : n) + logsqdev .- _log_count(logsqdev, max(0, corrected ? n - 1 : n)) # Single pass over a general iterator: (logsumexp(X), count). function _logsumexp_and_count(X) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 05c04c84..eab35cf1 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,5 +1,6 @@ using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var +using OffsetArrays: OffsetArray using LogExpFunctions: logmeanexp, logstdexp, logvarexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp @@ -7,9 +8,6 @@ using LogExpFunctions: logmeanexp, logstdexp, logvarexp, # arguments so that they are concretely typed inside this function (avoiding spurious # allocations from captured, boxed locals). allocations(f, x) = (f(x); @allocated f(x)) -# Same, for the `dims` keyword form (`x` is preallocated by the caller, so only the -# allocations of `f(x; dims=1)` itself are counted). -allocations_dims(f, x) = (f(x; dims=1); @allocated f(x; dims=1)) @testset "logmeanexp, logvarexp, logstdexp arrays" begin for T in (Float32, Float64) @@ -81,6 +79,57 @@ end @test isnan(logstdexp(Xsingleton; dims=:, corrected=true)) end +# Regressions for correctness bugs found in review. Each block is one bug class. +@testset "edge-case regressions" begin + # Non-1-based axes (OffsetArrays): the `dims` variance/std must match the result on + # the equivalent 1-based array (a fused lazy reduction silently returned wrong values). + base = randn(4, 3) + oa = OffsetArray(base, -1, -1) + for dims in (1, 2, :) + @test collect(logmeanexp(oa; dims=dims)) ≈ collect(logmeanexp(base; dims=dims)) + for corrected in (true, false) + @test collect(logvarexp(oa; dims=dims, corrected=corrected)) ≈ + collect(logvarexp(base; dims=dims, corrected=corrected)) + @test collect(logstdexp(oa; dims=dims, corrected=corrected)) ≈ + collect(logstdexp(base; dims=dims, corrected=corrected)) + end + end + + # Abstract element type: the `dims` variance must still work (and match a concretely + # typed copy), not throw a 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 logmeanexp_and_logvarexp(Xabstract; dims=2)[2] ≈ logvarexp(Xconcrete; dims=2) + + # Empty reduction along the reduced dimension: `var` is NaN (not an error) for every + # `corrected`, matching `Statistics.var` and `logmeanexp`. + 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=corrected)) + @test all(isnan, logstdexp(Eredux; dims=1, corrected=corrected)) + @test all(isnan, logmeanexp_and_logvarexp(Eredux; dims=1, corrected=corrected)[2]) + end + + # Empty along a dimension that is NOT being reduced: the result is an 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) + + # Complex arrays are rejected with a clear ArgumentError on every variance/std path, + # with or without `dims` (previously the `dims` form threw a confusing MethodError). + C = ComplexF64[1 2; 3 4] + @test_throws ArgumentError logvarexp(C) + @test_throws ArgumentError logvarexp(C; dims=1) + @test_throws ArgumentError logstdexp(C; dims=2) + @test_throws ArgumentError logmeanexp_and_logvarexp(C; dims=1) + @test_throws ArgumentError logmeanexp_and_logstdexp(C; dims=1) +end + @testset "logmeanexp_and_logvarexp, logmeanexp_and_logstdexp" begin for T in (Float32, Float64) X = randn(T, 5, 3, 2) @@ -140,13 +189,6 @@ end # genuinely allocation-free paths @test allocations(logmeanexp, randn(T, 10_000)) == 0 @test allocations(logvarexp, Tuple(randn(T, 20))) == 0 - # `dims` reductions allocate only the (fixed-size) reduced output, never a - # temporary the size of the input: the variance is fused into `logsumexp` lazily. - # Reducing along `dims=1` gives the same `(1, 4)` output for both sizes, so a - # constant allocation count confirms no O(length(X)) intermediate is built. - for f in (logvarexp, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp) - @test allocations_dims(f, randn(T, 10_000, 4)) == allocations_dims(f, randn(T, 100, 4)) - end end end From 695b2efc9397f460098a4216265463e6772c5352 Mon Sep 17 00:00:00 2001 From: cossio Date: Thu, 4 Jun 2026 23:17:43 +0200 Subject: [PATCH 16/40] Simplify: dedup the one-pass reduction loops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_centered_logsqdev` and `_logsumexp_and_count` were two near-identical hand-rolled iterate/accumulate loops differing only by a per-element transform. Replace both with a single `_logsumexp_count(f, X)` that returns `(logsumexp(f(xᵢ)), count)` in one pass — `logmeanexp` uses `f = identity`, `_centered_logsqdev` uses the centered-square term via `Base.Fix2`. Also factor the duplicated "require real inputs" message into `_throw_not_real` so the two validators can't drift. No behavior change. Verified on Julia 1.10.11 and 1.12.6: allocations, type-stability/inference, JET and Aqua all unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 40 ++++++++++++++++------------------------ 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index c487c5b7..07c9e18b 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -21,7 +21,7 @@ function logmeanexp(X) n = _known_length(X) if isnothing(n) # length not known ahead of time: count during the single pass - lse, n = _logsumexp_and_count(X) + lse, n = _logsumexp_count(identity, X) else iszero(n) && _throw_empty() lse = logsumexp(X) @@ -146,14 +146,16 @@ _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) # count is irrelevant — the result is empty — so avoid dividing by zero. _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) -_require_real(x::Real) = x -_require_real(x) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) +_throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) -# variance/std require real inputs; reject a non-real element type up front (with the same -# message as `_require_real`) so the array `dims` paths fail cleanly instead of hitting a -# `MethodError` deep inside `logsubexp`. +# variance/std require real inputs. `_require_real` checks a single value (the generic +# iterator path, whose element type may be unknown); `_require_real_array` rejects a +# non-real element type up front so the array `dims` paths fail cleanly with the same +# message instead of hitting a `MethodError` deep inside `logsubexp`. +_require_real(x::Real) = x +_require_real(x) = _throw_not_real() _require_real_array(X::AbstractArray{<:Real}) = X -_require_real_array(X::AbstractArray) = throw(ArgumentError("logvarexp and logstdexp require real inputs")) +_require_real_array(X::AbstractArray) = _throw_not_real() _known_length(X) = _known_length(Base.IteratorSize(typeof(X)), X) _known_length(::Union{Base.HasLength,Base.HasShape}, X) = length(X) @@ -173,19 +175,7 @@ _logsqdev_term(x, logmean) = 2 * logsubexp(_require_real(x), logmean) # log of the sum of squared deviations, `logsumexp(2 * logsubexp(xᵢ, logmean))`, # accumulated in a single pass without allocating an intermediate array. -function _centered_logsqdev(X, logmean) - next = iterate(X) - isnothing(next) && _throw_empty() - x, state = next - acc = _logsqdev_term(x, logmean) - while true - next = iterate(X, state) - isnothing(next) && break - x, state = next - acc = _logsumexp_onepass_op(acc, _logsqdev_term(x, logmean)) - end - return _logsumexp_onepass_result(acc) -end +_centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) # divide the squared-deviation sum by the count, in log space. When the (corrected) count # is non-positive — a single element with `corrected=true`, or an empty reduction — the @@ -194,18 +184,20 @@ end _finish_logvar(logsqdev, n::Integer, corrected::Bool) = logsqdev .- _log_count(logsqdev, max(0, corrected ? n - 1 : n)) -# Single pass over a general iterator: (logsumexp(X), count). -function _logsumexp_and_count(X) +# One pass over a general iterator, returning `(logsumexp(f(xᵢ)), count)`. Used both to +# average (`f = identity`, in `logmeanexp`) and to sum squared deviations (`f` the centered +# square term, in `_centered_logsqdev`). +function _logsumexp_count(f, X) next = iterate(X) isnothing(next) && _throw_empty() x, state = next - acc = x + acc = f(x) n = 1 while true next = iterate(X, state) isnothing(next) && break x, state = next - acc = _logsumexp_onepass_op(acc, x) + acc = _logsumexp_onepass_op(acc, f(x)) n += 1 end return _logsumexp_onepass_result(acc), n From 23f5dc469e05b116e0d01da3c251fbc9bc6c8a37 Mon Sep 17 00:00:00 2001 From: cossio Date: Fri, 5 Jun 2026 00:17:48 +0200 Subject: [PATCH 17/40] Make logmeanexp single-pass; fix single-use HasLength iterators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generic logmeanexp(X) had a length-known branch that called logsumexp(X) (which probes isempty(X)) plus a separate length(X). For a single-use iterator that reports a length, the isempty probe consumed the first element, so the mean was computed over the remaining elements but divided by the full count — a silently wrong result (the variance path already materialized such iterators, so mean and variance were inconsistent). Always count during the single _logsumexp_count pass instead. This also drops the two-branch structure and the `_known_length` helper (counting requires a full pass regardless, so the length shortcut never saved work). The array method is unaffected (arrays dispatch to their own method). Adds a regression test with a single-use, length-reporting iterator. Verified on Julia 1.10.11 and 1.12.6 (allocations, inference, JET, Aqua green). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 15 +++------------ test/logstatsexp.jl | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 07c9e18b..99206110 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -18,14 +18,9 @@ Compute `log(mean(exp, X))` in a numerically stable way. `X` should be an iterator of numbers. """ function logmeanexp(X) - n = _known_length(X) - if isnothing(n) - # length not known ahead of time: count during the single pass - lse, n = _logsumexp_count(identity, X) - else - iszero(n) && _throw_empty() - lse = logsumexp(X) - end + # single pass, counting as we go: never traverses `X` twice, so single-use iterators + # (and ones that only report a length) are handled correctly. + lse, n = _logsumexp_count(identity, X) return lse - _log_count(lse, n) end @@ -157,10 +152,6 @@ _require_real(x) = _throw_not_real() _require_real_array(X::AbstractArray{<:Real}) = X _require_real_array(X::AbstractArray) = _throw_not_real() -_known_length(X) = _known_length(Base.IteratorSize(typeof(X)), X) -_known_length(::Union{Base.HasLength,Base.HasShape}, X) = length(X) -_known_length(_, X) = nothing - # Variance is centered, so we need to traverse the data twice (mean, then deviations). # Known re-iterable containers are traversed in place; any other iterator is materialized # once, so that one-shot iterators (e.g. `Iterators.Stateful`) are handled correctly. diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index eab35cf1..e831d946 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -9,6 +9,23 @@ using LogExpFunctions: logmeanexp, logstdexp, logvarexp, # allocations from captured, boxed locals). allocations(f, x) = (f(x); @allocated f(x)) +# A single-use iterator that nonetheless advertises a length (so it exercises the +# length-aware code path). Iterating it consumes it; a second traversal — or an +# `isempty` probe — would skip elements. Used to check `logmeanexp`'s one-pass contract. +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) @@ -120,6 +137,11 @@ end @test size(logvarexp(Eother; dims=1)) == (1, 0) @test size(logstdexp(Eother; dims=1)) == (1, 0) + # Single-use iterator that reports a length: logmeanexp must traverse it exactly once + # (a length-aware path that re-consumed it would skip the first element). + data = randn(7) + @test logmeanexp(DrainOnce(data)) ≈ log(mean(exp, data)) + # Complex arrays are rejected with a clear ArgumentError on every variance/std path, # with or without `dims` (previously the `dims` form threw a confusing MethodError). C = ComplexF64[1 2; 3 4] From 8fafe93412e5fb19c24abf09a52a68bd7499635b Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 19:56:12 +0200 Subject: [PATCH 18/40] Apply suggestion from @devmotion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: David Müller-Widmann --- src/logstatsexp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 99206110..989ce98f 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -27,7 +27,7 @@ end """ $(SIGNATURES) -Compute `log.(mean(exp.(X); dims=dims))` in a numerically stable way. +Compute `log.(mean(exp.(X); dims))` in a numerically stable way. """ function logmeanexp(X::AbstractArray{<:Number}; dims=:) dims isa Colon && isempty(X) && _throw_empty() From c9439aa0c9c8a2b886a4a7a0db4e938919fd3211 Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 19:57:22 +0200 Subject: [PATCH 19/40] Apply suggestion from @devmotion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: David Müller-Widmann --- src/logstatsexp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 989ce98f..0114fba8 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -31,7 +31,7 @@ Compute `log.(mean(exp.(X); dims))` in a numerically stable way. """ function logmeanexp(X::AbstractArray{<:Number}; dims=:) dims isa Colon && isempty(X) && _throw_empty() - lse = logsumexp(X; dims=dims) + lse = logsumexp(X; dims) return lse .- _log_count(lse, _reduced_count(X, lse)) end From 1be77c1e76f7f643b25a624198e9dd018f239ffa Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 19:57:41 +0200 Subject: [PATCH 20/40] Apply suggestion from @devmotion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: David Müller-Widmann --- src/logstatsexp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 0114fba8..745729a2 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -38,7 +38,7 @@ end """ $(SIGNATURES) -Compute `(log(mean(exp, X)), log(var(exp, X; corrected=corrected)))` in a numerically +Compute `(log(mean(exp, X)), log(var(exp, X; corrected)))` in a numerically stable way. Computing the two together is cheaper than calling [`logmeanexp`](@ref) and [`logvarexp`](@ref) separately, since the mean is reused to center the variance. From 691e34e616db164fe36952b35b81f92968c7f230 Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 20:08:41 +0200 Subject: [PATCH 21/40] Clean up logstatsexp per review feedback Address review comments on JuliaStats#120 (no behavior change except empty array mean, see below): - Trim the verbose/editorializing header and helper comments; keep only the formula and the non-obvious rationale. - Merge the duplicated per-method docstrings into one docstring per function that also documents the array `dims` form. - Scalar `logmeanexp` now computes `lse - log(oftype(lse, n))` directly instead of going through the `_log_count` helper. - `logmeanexp(::AbstractArray; dims)` subtracts in place for the array case, avoiding a second temporary; empty-array reduction now returns `NaN` (matching `mean`) instead of throwing. - Use keyword-argument punning (`; dims`, `; corrected`) throughout. `_materialize` is kept (it lets re-iterable containers stay allocation-free while collecting single-use iterators for the variance's two passes). Tested on Julia 1.10 and 1.12: full logstatsexp testset, JET, and Aqua pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/logstatsexp.jl | 133 +++++++++++++++----------------------------- test/logstatsexp.jl | 2 +- 2 files changed, 45 insertions(+), 90 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 745729a2..e1cd60c5 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -1,129 +1,100 @@ # Numerically stable `log`-of-statistics-of-`exp` reductions. # -# The mean is `logsumexp(X) - log(n)`. The variance uses the *centered* formula +# The mean is `logsumexp(X) - log(n)`. The variance uses the centered formula # # log var = logsumexp(2 * logsubexp(xᵢ, logmean)) - log(n - corrected) # -# i.e. the log of the sum of squared deviations `∑ᵢ (exp(xᵢ) - mean)²`, divided by the -# count. Centering is essential for numerical stability: the raw-moment alternative -# `logsubexp(∑exp(2xᵢ), (∑exp(xᵢ))²/n)` cancels catastrophically when the variance is -# small relative to the mean (and can even overflow to `Inf` for nearly-equal inputs). -# Computing the variance therefore needs the mean first, hence the two helpers below. +# i.e. the log of the sum of squared deviations `∑ᵢ (exp(xᵢ) - mean)²`, divided by the count. """ $(SIGNATURES) Compute `log(mean(exp, X))` in a numerically stable way. -`X` should be an iterator of numbers. +`X` should be an iterator of numbers. For an array, `dims` selects the dimensions to reduce +over, returning `log.(mean(exp.(X); dims))`. """ function logmeanexp(X) - # single pass, counting as we go: never traverses `X` twice, so single-use iterators - # (and ones that only report a length) are handled correctly. lse, n = _logsumexp_count(identity, X) - return lse - _log_count(lse, n) + return lse - log(oftype(lse, n)) end -""" -$(SIGNATURES) - -Compute `log.(mean(exp.(X); dims))` in a numerically stable way. -""" function logmeanexp(X::AbstractArray{<:Number}; dims=:) - dims isa Colon && isempty(X) && _throw_empty() lse = logsumexp(X; dims) - return lse .- _log_count(lse, _reduced_count(X, lse)) + c = _log_count(lse, _reduced_count(X, lse)) + lse isa Number && return lse - c + lse .-= c + return lse end """ $(SIGNATURES) -Compute `(log(mean(exp, X)), log(var(exp, X; corrected)))` in a numerically -stable way. Computing the two together is cheaper than calling [`logmeanexp`](@ref) and +Compute `(log(mean(exp, X)), log(var(exp, X; corrected)))` in a numerically stable way. +Computing the two together is cheaper than calling [`logmeanexp`](@ref) and [`logvarexp`](@ref) separately, since the mean is reused to center the variance. -`X` should be an iterator of real numbers. +`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to +reduce over, returning `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))`. """ function logmeanexp_and_logvarexp(X; corrected::Bool=true) - xs = _materialize(X) + xs = _materialize(X) # variance needs two passes (mean, then deviations) logmean = logmeanexp(xs) logsqdev = _centered_logsqdev(xs, logmean) return logmean, _finish_logvar(logsqdev, length(xs), corrected) end -""" -$(SIGNATURES) - -Compute `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))` in a -numerically stable way, reusing the mean to center the variance. -""" function logmeanexp_and_logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) _require_real_array(X) - logmean = logmeanexp(X; dims=dims) + logmean = logmeanexp(X; dims) return logmean, _centered_logvar(X, logmean, dims, corrected) end # dispatch on `dims` so the return type is concrete (no `Union` of scalar/array results) -_centered_logvar(X, logmean, ::Colon, corrected::Bool) = # scalar reduction, no temporary +_centered_logvar(X, logmean, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) _centered_logvar(X, logmean, dims, corrected::Bool) = - _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims=dims), _reduced_count(X, logmean), corrected) + _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims), _reduced_count(X, logmean), corrected) """ $(SIGNATURES) -Compute `log(var(exp, X; corrected=corrected))` in a numerically stable way. - -`X` should be an iterator of real numbers. -""" -logvarexp(X; corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; corrected=corrected)) - -""" -$(SIGNATURES) +Compute `log(var(exp, X; corrected))` in a numerically stable way. -Compute `log.(var(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. +`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to +reduce over, returning `log.(var(exp.(X); dims, corrected))`. """ +logvarexp(X; corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; corrected)) logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - last(logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected)) + last(logmeanexp_and_logvarexp(X; dims, corrected)) """ $(SIGNATURES) -Compute `log(std(exp, X; corrected=corrected))` in a numerically stable way. - -`X` should be an iterator of real numbers. -""" -logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected=corrected) / 2 - -""" -$(SIGNATURES) +Compute `log(std(exp, X; corrected))` in a numerically stable way. -Compute `log.(std(exp.(X); dims=dims, corrected=corrected))` in a numerically stable way. +`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to +reduce over, returning `log.(std(exp.(X); dims, corrected))`. """ +logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected) / 2 logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - logvarexp(X; dims=dims, corrected=corrected) / 2 + logvarexp(X; dims, corrected) / 2 """ $(SIGNATURES) -Compute `(log(mean(exp, X)), log(std(exp, X; corrected=corrected)))` in a numerically -stable way, reusing the mean to center the variance. +Compute `(log(mean(exp, X)), log(std(exp, X; corrected)))` in a numerically stable way, +reusing the mean to center the variance. -`X` should be an iterator of real numbers. +`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to +reduce over, returning `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))`. """ function logmeanexp_and_logstdexp(X; corrected::Bool=true) - logmean, logvar = logmeanexp_and_logvarexp(X; corrected=corrected) + logmean, logvar = logmeanexp_and_logvarexp(X; corrected) return logmean, logvar / 2 end - -""" -$(SIGNATURES) - -Compute `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))` in a -numerically stable way, reusing the mean to center the variance. -""" function logmeanexp_and_logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) - logmean, logvar = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) + logmean, logvar = logmeanexp_and_logvarexp(X; dims, corrected) return logmean, logvar / 2 end @@ -131,53 +102,37 @@ end _throw_empty() = throw(ArgumentError("reducing over an empty collection is not allowed")) -# `log(n)` in the (real) floating point type of `R`, so that no unwanted promotion of -# `R` (e.g. `Float32` to `Float64`) takes place. Works for scalar and array `R` alike. -# `n == 0` gives `log(0) == -Inf`, which produces the expected `NaN`/`-Inf` results. +# `log(n)` in the (real) float type of `R`, avoiding promotion (e.g. `Float32` to `Float64`). +# Works for scalar and array `R`; `n == 0` gives `-Inf`, yielding the expected `NaN`/`-Inf`. _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) -# number of elements reduced into each entry of `R`. When `R` is empty (the reduction -# produced no cells, e.g. `X` is empty along a dimension that is not being reduced) the -# count is irrelevant — the result is empty — so avoid dividing by zero. +# number of elements reduced into each entry of `R` (empty `R` ⇒ 0, avoiding division by zero) _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) +# variance/std require real inputs; reject a non-real element type up front for a clean error _throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) - -# variance/std require real inputs. `_require_real` checks a single value (the generic -# iterator path, whose element type may be unknown); `_require_real_array` rejects a -# non-real element type up front so the array `dims` paths fail cleanly with the same -# message instead of hitting a `MethodError` deep inside `logsubexp`. _require_real(x::Real) = x _require_real(x) = _throw_not_real() _require_real_array(X::AbstractArray{<:Real}) = X _require_real_array(X::AbstractArray) = _throw_not_real() -# Variance is centered, so we need to traverse the data twice (mean, then deviations). -# Known re-iterable containers are traversed in place; any other iterator is materialized -# once, so that one-shot iterators (e.g. `Iterators.Stateful`) are handled correctly. -# (`IteratorSize` cannot be used to detect re-iterability: `Stateful` reports `HasLength` -# on some Julia versions yet is single-use.) +# re-iterable containers are traversed in place; any other iterator is collected once, so that +# single-use iterators (e.g. `Iterators.Stateful`) survive the variance's two passes _materialize(X) = collect(X) _materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X -# log of the (centered) squared deviation of a single point, `2 * logsubexp(xᵢ, logmean)` -# (i.e. `log((exp(xᵢ) - mean)^2)`). This is the term summed to form the variance. +# `log((exp(xᵢ) - mean)^2) = 2 * logsubexp(xᵢ, logmean)`, the term summed for the variance _logsqdev_term(x, logmean) = 2 * logsubexp(_require_real(x), logmean) -# log of the sum of squared deviations, `logsumexp(2 * logsubexp(xᵢ, logmean))`, -# accumulated in a single pass without allocating an intermediate array. +# `logsumexp(2 * logsubexp(xᵢ, logmean))`, accumulated in a single pass _centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) -# divide the squared-deviation sum by the count, in log space. When the (corrected) count -# is non-positive — a single element with `corrected=true`, or an empty reduction — the -# numerator and `log` of the (clamped-to-zero) denominator are both `-Inf`, giving `NaN`, -# matching `var`. Clamping to zero also avoids `log(-1)` for an empty `corrected=true` case. +# divide the squared-deviation sum by the count in log space; a non-positive count (a single +# element with `corrected=true`, or an empty reduction) gives `NaN`, matching `var` _finish_logvar(logsqdev, n::Integer, corrected::Bool) = logsqdev .- _log_count(logsqdev, max(0, corrected ? n - 1 : n)) -# One pass over a general iterator, returning `(logsumexp(f(xᵢ)), count)`. Used both to -# average (`f = identity`, in `logmeanexp`) and to sum squared deviations (`f` the centered -# square term, in `_centered_logsqdev`). +# one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` function _logsumexp_count(f, X) next = iterate(X) isnothing(next) && _throw_empty() diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index e831d946..08033d9d 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -70,7 +70,7 @@ end @test_throws ArgumentError logvarexp((1.0 + 0.0im, 2.0 + 0.0im)) @test_throws ArgumentError logstdexp((1.0 + 0.0im, 2.0 + 0.0im)) @test_throws ArgumentError logmeanexp(()) - @test_throws ArgumentError logmeanexp(Float64[]) + @test isnan(logmeanexp(Float64[])) # empty array mean is NaN, matching `mean` @test_throws ArgumentError logvarexp(()) @test_throws ArgumentError logvarexp(Float64[]) end From 09a8da12654e38d8d1e7a1331d3fae9604e9b0dc Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 21:47:03 +0200 Subject: [PATCH 22/40] Add in-place logmeanexp!/logvarexp!; dims reductions allocate only the output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the structure of `src/logsumexp.jl`: - Add exported `logmeanexp!(out, X)` and `logvarexp!(out, X; corrected)` that reduce over the singleton dimensions of `out` and write the result there, reusing `logsumexp!` (so they allocate only the same `Tuple{FT,FT}` scratch `logsumexp!` itself uses — no input-sized temporary). - `logmeanexp`/`logvarexp`/`logstdexp`/`logmeanexp_and_log*` with `dims` now allocate the (output-sized) result and delegate to the in-place versions, instead of materializing the full `2 .* logsubexp.(X, logmean)` array. The `dims` variance went from O(n) to O(output) allocations. - The centered terms are reduced lazily through a small `_LogSqDev` array (a plain `AbstractArray`, not a `Broadcasted`, so `reduce`/`reducedim!` honor offset axes and a concrete element type). Tests: add an in-place correctness testset (dense, OffsetArray, abstract eltype, complex rejection) and extend the allocation tests to assert the `dims` and in-place paths do not scale with the input size. Verified on Julia 1.10 and 1.12: full logstatsexp testset, JET, and Aqua pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/index.md | 2 + src/LogExpFunctions.jl | 2 +- src/logstatsexp.jl | 90 +++++++++++++++++++++++++++++++++--------- test/logstatsexp.jl | 31 ++++++++++++++- 4 files changed, 104 insertions(+), 21 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index d2d33fbe..c51b66c2 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -39,7 +39,9 @@ logitexp log1mlogistic logit1mexp logmeanexp +logmeanexp! logvarexp +logvarexp! logstdexp logmeanexp_and_logvarexp logmeanexp_and_logstdexp diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index 3d8cac3a..ebe46b8b 100644 --- a/src/LogExpFunctions.jl +++ b/src/LogExpFunctions.jl @@ -9,7 +9,7 @@ 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, logvarexp, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp, + logmeanexp, logmeanexp!, logvarexp, logvarexp!, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp, loglogistic, logitexp, log1mlogistic, logit1mexp include("basicfuns.jl") diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index e1cd60c5..36a45961 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -13,20 +13,30 @@ 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) lse, n = _logsumexp_count(identity, X) return lse - log(oftype(lse, n)) end +logmeanexp(X::AbstractArray{<:Number}; dims=:) = _logmeanexp(X, dims) + +""" +$(SIGNATURES) -function logmeanexp(X::AbstractArray{<:Number}; dims=:) - lse = logsumexp(X; dims) - c = _log_count(lse, _reduced_count(X, lse)) - lse isa Number && return lse - c - lse .-= c - return lse +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) + out .-= _log_count(out, _reduced_count(X, out)) + return out end +_logmeanexp(X, ::Colon) = (lse = logsumexp(X); lse - _log_count(lse, length(X))) +_logmeanexp(X, dims) = logmeanexp!(_reduced_similar(X, dims), X) + """ $(SIGNATURES) @@ -38,23 +48,21 @@ Computing the two together is cheaper than calling [`logmeanexp`](@ref) and reduce over, returning `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))`. """ function logmeanexp_and_logvarexp(X; corrected::Bool=true) - xs = _materialize(X) # variance needs two passes (mean, then deviations) + xs = _materialize(X) logmean = logmeanexp(xs) logsqdev = _centered_logsqdev(xs, logmean) return logmean, _finish_logvar(logsqdev, length(xs), corrected) end - function logmeanexp_and_logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) _require_real_array(X) logmean = logmeanexp(X; dims) return logmean, _centered_logvar(X, logmean, dims, corrected) end -# dispatch on `dims` so the return type is concrete (no `Union` of scalar/array results) _centered_logvar(X, logmean, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) _centered_logvar(X, logmean, dims, corrected::Bool) = - _finish_logvar(logsumexp(2 .* logsubexp.(X, logmean); dims), _reduced_count(X, logmean), corrected) + _finish_logvar(logsumexp(_LogSqDev(X, logmean); dims), _reduced_count(X, logmean), corrected) """ $(SIGNATURES) @@ -63,10 +71,31 @@ Compute `log(var(exp, X; corrected))` in a numerically stable way. `X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to reduce over, returning `log.(var(exp.(X); dims, corrected))`. + +See also [`logvarexp!`](@ref). """ logvarexp(X; corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; corrected)) logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - last(logmeanexp_and_logvarexp(X; dims, corrected)) + _logvarexp(_require_real_array(X), dims, corrected) + +""" +$(SIGNATURES) + +Compute [`logvarexp`](@ref) of `X` over the singleton dimensions of `out`, and write the +result to `out`. +""" +function logvarexp!(out::AbstractArray, X::AbstractArray{<:Number}; corrected::Bool=true) + _require_real_array(X) + logmeanexp!(out, X) + n = _reduced_count(X, out) + logsumexp!(out, _LogSqDev(X, out)) + out .-= _log_count(out, max(0, corrected ? n - 1 : n)) + return out +end + +_logvarexp(X, ::Colon, corrected::Bool) = + _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) +_logvarexp(X, dims, corrected::Bool) = logvarexp!(_reduced_similar(X, dims), X; corrected) """ $(SIGNATURES) @@ -78,7 +107,7 @@ reduce over, returning `log.(std(exp.(X); dims, corrected))`. """ logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected) / 2 logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - logvarexp(X; dims, corrected) / 2 + _halve(logvarexp(X; dims, corrected)) """ $(SIGNATURES) @@ -95,7 +124,7 @@ function logmeanexp_and_logstdexp(X; corrected::Bool=true) end function logmeanexp_and_logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) logmean, logvar = logmeanexp_and_logvarexp(X; dims, corrected) - return logmean, logvar / 2 + return logmean, _halve(logvar) end # ---- internal helpers ---- @@ -109,6 +138,12 @@ _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) # number of elements reduced into each entry of `R` (empty `R` ⇒ 0, avoiding division by zero) _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) +# output array for a reduction of `X` over `dims` (singleton along the reduced dimensions) +_reduced_similar(X, dims) = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) + +_halve(v::Number) = v / 2 +_halve(v::AbstractArray) = v ./= 2 + # variance/std require real inputs; reject a non-real element type up front for a clean error _throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) _require_real(x::Real) = x @@ -123,14 +158,14 @@ _materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X # `log((exp(xᵢ) - mean)^2) = 2 * logsubexp(xᵢ, logmean)`, the term summed for the variance _logsqdev_term(x, logmean) = 2 * logsubexp(_require_real(x), logmean) - -# `logsumexp(2 * logsubexp(xᵢ, logmean))`, accumulated in a single pass _centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) -# divide the squared-deviation sum by the count in log space; a non-positive count (a single -# element with `corrected=true`, or an empty reduction) gives `NaN`, matching `var` -_finish_logvar(logsqdev, n::Integer, corrected::Bool) = - logsqdev .- _log_count(logsqdev, max(0, corrected ? n - 1 : n)) +function _finish_logvar(logsqdev, n::Integer, corrected::Bool) + c = _log_count(logsqdev, max(0, corrected ? n - 1 : n)) + logsqdev isa Number && return logsqdev - c + logsqdev .-= c + return logsqdev +end # one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` function _logsumexp_count(f, X) @@ -148,3 +183,20 @@ function _logsumexp_count(f, X) end return _logsumexp_onepass_result(acc), n end + +# lazy array of the centered terms `2 * logsubexp(xᵢ, meanⱼ)`, with `mean` broadcast over the +# reduced (singleton) dimensions. A plain `AbstractArray` (not a `Broadcasted`) so that +# `reduce`/`reducedim!` honor offset axes and a concrete element type. +struct _LogSqDev{T,N,XT<:AbstractArray,MT<:AbstractArray} <: AbstractArray{T,N} + x::XT + mean::MT +end +_LogSqDev(x::AbstractArray, mean::AbstractArray) = + _LogSqDev{float(eltype(x)),ndims(x),typeof(x),typeof(mean)}(x, mean) +Base.size(s::_LogSqDev) = size(s.x) +Base.axes(s::_LogSqDev) = axes(s.x) +Base.IndexStyle(::Type{<:_LogSqDev}) = IndexCartesian() +Base.@propagate_inbounds function Base.getindex(s::_LogSqDev{T,N}, I::Vararg{Int,N}) where {T,N} + j = map((i, ax) -> ifelse(length(ax) == 1, first(ax), i), I, axes(s.mean)) + return convert(T, _logsqdev_term(s.x[I...], s.mean[j...])) +end diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 08033d9d..1b0015db 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,7 +1,7 @@ using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var using OffsetArrays: OffsetArray -using LogExpFunctions: logmeanexp, logstdexp, logvarexp, +using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logvarexp, logvarexp!, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp # Count heap allocations of `f(x)` after warming it up. `f` and `x` are passed as @@ -200,6 +200,25 @@ end @test typeof(@inferred(logmeanexp_and_logvarexp(X))) == Tuple{Float32,Float32} end +@testset "in-place logmeanexp!/logvarexp!" 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=dims) + for corrected in (true, false) + outv = similar(A, T, Base.reduced_indices(axes(A), dims)) + @test logvarexp!(outv, A; corrected=corrected) === outv + @test outv ≈ logvarexp(A; dims=dims, corrected=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_throws ArgumentError logvarexp!(Matrix{ComplexF64}(undef, 1, 2), ComplexF64[1 2; 3 4]) +end + @testset "allocations" begin # Full reductions allocate at most a small constant (no per-element / O(n) temporary): # the allocation count must not grow with the input size. @@ -207,10 +226,20 @@ end for f in (logmeanexp, logvarexp, logstdexp, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp) @test allocations(f, randn(T, 10_000)) == allocations(f, randn(T, 100)) + # `dims` reductions only allocate the (output-sized) result and scratch — never an + # O(n) temporary — so allocations are independent of the reduced dimension's length. + g = X -> f(X; dims=1) + @test allocations(g, randn(T, 10_000, 3)) == allocations(g, randn(T, 100, 3)) end # genuinely allocation-free paths @test allocations(logmeanexp, randn(T, 10_000)) == 0 @test allocations(logvarexp, Tuple(randn(T, 20))) == 0 + # in-place reductions reuse `out`, so their allocation does not grow with the input + 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)) + @test allocations(X -> logvarexp!(out, X), randn(T, 10_000, 3)) == + allocations(X -> logvarexp!(out, X), randn(T, 100, 3)) end end From 5601c23a882b920a233ccbcf4671503f286bdffc Mon Sep 17 00:00:00 2001 From: cossio Date: Mon, 8 Jun 2026 22:28:20 +0200 Subject: [PATCH 23/40] Use keyword-argument punning in logstatsexp tests Co-Authored-By: Claude Opus 4.8 (1M context) --- test/logstatsexp.jl | 60 ++++++++++++++++++++++----------------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 1b0015db..56b2d820 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -30,12 +30,12 @@ end for T in (Float32, Float64) X = randn(T, 5, 3, 2) for dims in (2, (1, 2), :) - @test logmeanexp(X; dims=dims) ≈ log.(mean(exp.(X); dims=dims)) + @test logmeanexp(X; dims) ≈ log.(mean(exp.(X); dims)) for corrected in (true, false) - @test logvarexp(X; dims=dims, corrected=corrected) ≈ - log.(var(exp.(X); dims=dims, corrected=corrected)) - @test logstdexp(X; dims=dims, corrected=corrected) ≈ - log.(std(exp.(X); dims=dims, corrected=corrected)) + @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)) @@ -79,9 +79,9 @@ end X = randn(Float32, 5, 3, 2) for dims in (1, (2, 3)) - @test eltype(@inferred(logmeanexp(X; dims=dims))) == Float32 - @test eltype(@inferred(logvarexp(X; dims=dims))) == Float32 - @test eltype(@inferred(logstdexp(X; dims=dims))) == Float32 + @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 @@ -103,12 +103,12 @@ end base = randn(4, 3) oa = OffsetArray(base, -1, -1) for dims in (1, 2, :) - @test collect(logmeanexp(oa; dims=dims)) ≈ collect(logmeanexp(base; dims=dims)) + @test collect(logmeanexp(oa; dims)) ≈ collect(logmeanexp(base; dims)) for corrected in (true, false) - @test collect(logvarexp(oa; dims=dims, corrected=corrected)) ≈ - collect(logvarexp(base; dims=dims, corrected=corrected)) - @test collect(logstdexp(oa; dims=dims, corrected=corrected)) ≈ - collect(logstdexp(base; dims=dims, corrected=corrected)) + @test collect(logvarexp(oa; dims, corrected)) ≈ + collect(logvarexp(base; dims, corrected)) + @test collect(logstdexp(oa; dims, corrected)) ≈ + collect(logstdexp(base; dims, corrected)) end end @@ -125,9 +125,9 @@ end 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=corrected)) - @test all(isnan, logstdexp(Eredux; dims=1, corrected=corrected)) - @test all(isnan, logmeanexp_and_logvarexp(Eredux; dims=1, corrected=corrected)[2]) + @test all(isnan, logvarexp(Eredux; dims=1, corrected)) + @test all(isnan, logstdexp(Eredux; dims=1, corrected)) + @test all(isnan, logmeanexp_and_logvarexp(Eredux; dims=1, corrected)[2]) end # Empty along a dimension that is NOT being reduced: the result is an empty array of @@ -156,12 +156,12 @@ end for T in (Float32, Float64) X = randn(T, 5, 3, 2) for dims in (2, (1, 2), :), corrected in (true, false) - m, v = logmeanexp_and_logvarexp(X; dims=dims, corrected=corrected) - @test m ≈ logmeanexp(X; dims=dims) - @test v ≈ logvarexp(X; dims=dims, corrected=corrected) - m2, s = logmeanexp_and_logstdexp(X; dims=dims, corrected=corrected) - @test m2 ≈ logmeanexp(X; dims=dims) - @test s ≈ logstdexp(X; dims=dims, corrected=corrected) + m, v = logmeanexp_and_logvarexp(X; dims, corrected) + @test m ≈ logmeanexp(X; dims) + @test v ≈ logvarexp(X; dims, corrected) + m2, s = logmeanexp_and_logstdexp(X; dims, corrected) + @test m2 ≈ logmeanexp(X; dims) + @test s ≈ logstdexp(X; dims, corrected) end # results match the reference statistics directly @test all(logmeanexp_and_logvarexp(X) .≈ (log(mean(exp, X)), log(var(exp.(X))))) @@ -188,8 +188,8 @@ end X = randn(Float32, 5, 3, 2) xt = Tuple(randn(Float32, 20)) for dims in (1, (2, 3), :) - @test @inferred(logmeanexp_and_logvarexp(X; dims=dims)) isa Tuple - @test @inferred(logmeanexp_and_logstdexp(X; dims=dims)) isa Tuple + @test @inferred(logmeanexp_and_logvarexp(X; dims)) isa Tuple + @test @inferred(logmeanexp_and_logstdexp(X; dims)) isa Tuple end @test @inferred(logmeanexp_and_logvarexp(xt)) isa NTuple{2,Float32} @test @inferred(logmeanexp_and_logstdexp(xt)) isa NTuple{2,Float32} @@ -206,11 +206,11 @@ end 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=dims) + @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=corrected) === outv - @test outv ≈ logvarexp(A; dims=dims, corrected=corrected) + @test logvarexp!(outv, A; corrected) === outv + @test outv ≈ logvarexp(A; dims, corrected) end end end @@ -249,7 +249,7 @@ end # values large/small enough that exp would over-/under-flow Float64. setprecision(BigFloat, 256) do refmean(x) = Float64(log(mean(exp.(big.(x))))) - refvar(x; corrected=true) = Float64(log(var(exp.(big.(x)); corrected=corrected))) + 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² @@ -264,8 +264,8 @@ end # an inherent cancellation makes the relative error meaningless @test logmeanexp(x) ≈ refmean(x) rtol = 1e-9 atol = 1e-10 for corrected in (true, false) - @test logvarexp(x; corrected=corrected) ≈ refvar(x; corrected=corrected) rtol = 1e-8 atol = 1e-9 - @test logstdexp(x; corrected=corrected) ≈ refvar(x; corrected=corrected) / 2 rtol = 1e-8 atol = 1e-9 + @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 # array and one-shot-iterator paths agree even in the cancellation regime From 9669d64023a0000b5b640ea1c5a7bac481b1f967 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 00:25:16 +0200 Subject: [PATCH 24/40] Remove combined logmeanexp_and_* functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop `logmeanexp_and_logvarexp` and `logmeanexp_and_logstdexp` (both the iterator and `dims` array methods), their exports, and their docs entries. They computed two statistics at once, which made them harder to review for little benefit over calling `logmeanexp` and `logvarexp`/`logstdexp`. - Reimplement the iterator `logvarexp(X)` standalone — it previously delegated to `logmeanexp_and_logvarexp`. It now materializes once, takes the mean, and reduces the centered squared deviations directly. `_materialize` is kept so single-use iterators survive the variance's two passes. - Remove the now-unused `_centered_logvar` helper (it only served the combined array method); `logvarexp!`/`_LogSqDev` are unchanged. - Prune the combined-function testsets and references. Verified on Julia 1.10 and 1.12: full logstatsexp testset, JET, and Aqua pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/index.md | 2 -- src/LogExpFunctions.jl | 2 +- src/logstatsexp.jl | 61 ++++++----------------------------------- test/logstatsexp.jl | 62 ++---------------------------------------- 4 files changed, 13 insertions(+), 114 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index c51b66c2..6c96ce3b 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -43,6 +43,4 @@ logmeanexp! logvarexp logvarexp! logstdexp -logmeanexp_and_logvarexp -logmeanexp_and_logstdexp ``` diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index ebe46b8b..19ff96a2 100644 --- a/src/LogExpFunctions.jl +++ b/src/LogExpFunctions.jl @@ -9,7 +9,7 @@ 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, logmeanexp_and_logvarexp, logmeanexp_and_logstdexp, + logmeanexp, logmeanexp!, logvarexp, logvarexp!, logstdexp, loglogistic, logitexp, log1mlogistic, logit1mexp include("basicfuns.jl") diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 36a45961..7a3f5c1f 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -40,33 +40,6 @@ _logmeanexp(X, dims) = logmeanexp!(_reduced_similar(X, dims), X) """ $(SIGNATURES) -Compute `(log(mean(exp, X)), log(var(exp, X; corrected)))` in a numerically stable way. -Computing the two together is cheaper than calling [`logmeanexp`](@ref) and -[`logvarexp`](@ref) separately, since the mean is reused to center the variance. - -`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to -reduce over, returning `(log.(mean(exp.(X); dims)), log.(var(exp.(X); dims, corrected)))`. -""" -function logmeanexp_and_logvarexp(X; corrected::Bool=true) - xs = _materialize(X) - logmean = logmeanexp(xs) - logsqdev = _centered_logsqdev(xs, logmean) - return logmean, _finish_logvar(logsqdev, length(xs), corrected) -end -function logmeanexp_and_logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) - _require_real_array(X) - logmean = logmeanexp(X; dims) - return logmean, _centered_logvar(X, logmean, dims, corrected) -end - -_centered_logvar(X, logmean, ::Colon, corrected::Bool) = - _finish_logvar(_centered_logsqdev(X, logmean), length(X), corrected) -_centered_logvar(X, logmean, dims, corrected::Bool) = - _finish_logvar(logsumexp(_LogSqDev(X, logmean); dims), _reduced_count(X, logmean), corrected) - -""" -$(SIGNATURES) - Compute `log(var(exp, X; corrected))` in a numerically stable way. `X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to @@ -74,7 +47,11 @@ reduce over, returning `log.(var(exp.(X); dims, corrected))`. See also [`logvarexp!`](@ref). """ -logvarexp(X; corrected::Bool=true) = last(logmeanexp_and_logvarexp(X; corrected)) +function logvarexp(X; corrected::Bool=true) + xs = _materialize(X) + logmean = logmeanexp(xs) + return _finish_logvar(_centered_logsqdev(xs, logmean), length(xs), corrected) +end logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = _logvarexp(_require_real_array(X), dims, corrected) @@ -107,30 +84,10 @@ reduce over, returning `log.(std(exp.(X); dims, corrected))`. """ logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected) / 2 logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - _halve(logvarexp(X; dims, corrected)) - -""" -$(SIGNATURES) - -Compute `(log(mean(exp, X)), log(std(exp, X; corrected)))` in a numerically stable way, -reusing the mean to center the variance. - -`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to -reduce over, returning `(log.(mean(exp.(X); dims)), log.(std(exp.(X); dims, corrected)))`. -""" -function logmeanexp_and_logstdexp(X; corrected::Bool=true) - logmean, logvar = logmeanexp_and_logvarexp(X; corrected) - return logmean, logvar / 2 -end -function logmeanexp_and_logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) - logmean, logvar = logmeanexp_and_logvarexp(X; dims, corrected) - return logmean, _halve(logvar) -end + _halve!(logvarexp(X; dims, corrected)) # ---- internal helpers ---- -_throw_empty() = throw(ArgumentError("reducing over an empty collection is not allowed")) - # `log(n)` in the (real) float type of `R`, avoiding promotion (e.g. `Float32` to `Float64`). # Works for scalar and array `R`; `n == 0` gives `-Inf`, yielding the expected `NaN`/`-Inf`. _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) @@ -141,8 +98,8 @@ _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) # output array for a reduction of `X` over `dims` (singleton along the reduced dimensions) _reduced_similar(X, dims) = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) -_halve(v::Number) = v / 2 -_halve(v::AbstractArray) = v ./= 2 +_halve!(v::Number) = v / 2 +_halve!(v::AbstractArray) = v ./= 2 # variance/std require real inputs; reject a non-real element type up front for a clean error _throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) @@ -170,7 +127,7 @@ end # one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` function _logsumexp_count(f, X) next = iterate(X) - isnothing(next) && _throw_empty() + isnothing(next) && throw(ArgumentError("reducing over an empty collection is not allowed")) x, state = next acc = f(x) n = 1 diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 56b2d820..4722bd0b 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,8 +1,7 @@ using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var using OffsetArrays: OffsetArray -using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logvarexp, logvarexp!, - logmeanexp_and_logvarexp, logmeanexp_and_logstdexp +using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logvarexp, logvarexp! # Count heap allocations of `f(x)` after warming it up. `f` and `x` are passed as # arguments so that they are concretely typed inside this function (avoiding spurious @@ -118,7 +117,7 @@ end Xconcrete = Float64.(Xabstract) @test logvarexp(Xabstract; dims=1) ≈ logvarexp(Xconcrete; dims=1) @test logvarexp(Xabstract) ≈ logvarexp(Xconcrete) - @test logmeanexp_and_logvarexp(Xabstract; dims=2)[2] ≈ logvarexp(Xconcrete; dims=2) + @test logvarexp(Xabstract; dims=2) ≈ logvarexp(Xconcrete; dims=2) # Empty reduction along the reduced dimension: `var` is NaN (not an error) for every # `corrected`, matching `Statistics.var` and `logmeanexp`. @@ -127,7 +126,6 @@ end for corrected in (true, false) @test all(isnan, logvarexp(Eredux; dims=1, corrected)) @test all(isnan, logstdexp(Eredux; dims=1, corrected)) - @test all(isnan, logmeanexp_and_logvarexp(Eredux; dims=1, corrected)[2]) end # Empty along a dimension that is NOT being reduced: the result is an empty array of @@ -148,56 +146,6 @@ end @test_throws ArgumentError logvarexp(C) @test_throws ArgumentError logvarexp(C; dims=1) @test_throws ArgumentError logstdexp(C; dims=2) - @test_throws ArgumentError logmeanexp_and_logvarexp(C; dims=1) - @test_throws ArgumentError logmeanexp_and_logstdexp(C; dims=1) -end - -@testset "logmeanexp_and_logvarexp, logmeanexp_and_logstdexp" begin - for T in (Float32, Float64) - X = randn(T, 5, 3, 2) - for dims in (2, (1, 2), :), corrected in (true, false) - m, v = logmeanexp_and_logvarexp(X; dims, corrected) - @test m ≈ logmeanexp(X; dims) - @test v ≈ logvarexp(X; dims, corrected) - m2, s = logmeanexp_and_logstdexp(X; dims, corrected) - @test m2 ≈ logmeanexp(X; dims) - @test s ≈ logstdexp(X; dims, corrected) - end - # results match the reference statistics directly - @test all(logmeanexp_and_logvarexp(X) .≈ (log(mean(exp, X)), log(var(exp.(X))))) - @test all(logmeanexp_and_logstdexp(X) .≈ (log(mean(exp, X)), log(std(exp.(X))))) - end - - # iterators (single pass, including one-shot iterators) - x = randn(Float32, 20) - xt = Tuple(x) - xe = exp.(x) - @test all(@inferred(logmeanexp_and_logvarexp(xt)) .≈ (log(mean(exp, xt)), log(var(xe)))) - @test all(@inferred(logmeanexp_and_logstdexp(xt)) .≈ (log(mean(exp, xt)), log(std(xe)))) - @test all(logmeanexp_and_logvarexp(Iterators.Stateful(x)) .≈ (log(mean(exp, x)), log(var(xe)))) - @test all(logmeanexp_and_logstdexp(Iterators.Stateful(x)) .≈ (log(mean(exp, x)), log(std(xe)))) - - # edge cases - @test isnan(last(logmeanexp_and_logvarexp((0.0,)))) - @test isnan(last(logmeanexp_and_logstdexp((0.0,)))) - @test_throws ArgumentError logmeanexp_and_logvarexp((1.0 + 0.0im,)) - @test_throws ArgumentError logmeanexp_and_logstdexp((1.0 + 0.0im,)) -end - -@testset "type stability and inference" begin - X = randn(Float32, 5, 3, 2) - xt = Tuple(randn(Float32, 20)) - for dims in (1, (2, 3), :) - @test @inferred(logmeanexp_and_logvarexp(X; dims)) isa Tuple - @test @inferred(logmeanexp_and_logstdexp(X; dims)) isa Tuple - end - @test @inferred(logmeanexp_and_logvarexp(xt)) isa NTuple{2,Float32} - @test @inferred(logmeanexp_and_logstdexp(xt)) isa NTuple{2,Float32} - - # no Float64 promotion for Float32 inputs - m, v = logmeanexp_and_logvarexp(X; dims=2) - @test eltype(m) == Float32 && eltype(v) == Float32 - @test typeof(@inferred(logmeanexp_and_logvarexp(X))) == Tuple{Float32,Float32} end @testset "in-place logmeanexp!/logvarexp!" begin @@ -223,8 +171,7 @@ end # Full reductions allocate at most a small constant (no per-element / O(n) temporary): # the allocation count must not grow with the input size. for T in (Float32, Float64) - for f in (logmeanexp, logvarexp, logstdexp, - logmeanexp_and_logvarexp, logmeanexp_and_logstdexp) + for f in (logmeanexp, logvarexp, logstdexp) @test allocations(f, randn(T, 10_000)) == allocations(f, randn(T, 100)) # `dims` reductions only allocate the (output-sized) result and scratch — never an # O(n) temporary — so allocations are independent of the reduced dimension's length. @@ -271,8 +218,5 @@ end # array and one-shot-iterator paths agree even in the cancellation regime x = 1.0 .+ 1e-5 .* randn(100) @test logvarexp(x) ≈ logvarexp(Iterators.Stateful(x)) rtol = 1e-10 - m, v = logmeanexp_and_logvarexp(Iterators.Stateful(x)) - @test m ≈ logmeanexp(x) rtol = 1e-10 - @test v ≈ logvarexp(x) rtol = 1e-10 end end From 8de5fea00d16acf301b1b133002056c9bdac8205 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 01:06:12 +0200 Subject: [PATCH 25/40] simplifications --- src/logstatsexp.jl | 35 +++++++++-------------------------- 1 file changed, 9 insertions(+), 26 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 7a3f5c1f..c4ca9d41 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -52,8 +52,8 @@ function logvarexp(X; corrected::Bool=true) logmean = logmeanexp(xs) return _finish_logvar(_centered_logsqdev(xs, logmean), length(xs), corrected) end -logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - _logvarexp(_require_real_array(X), dims, corrected) +logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = _logvarexp(X, dims, corrected) +logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = _throw_not_real() """ $(SIGNATURES) @@ -61,17 +61,16 @@ $(SIGNATURES) Compute [`logvarexp`](@ref) of `X` over the singleton dimensions of `out`, and write the result to `out`. """ -function logvarexp!(out::AbstractArray, X::AbstractArray{<:Number}; corrected::Bool=true) - _require_real_array(X) +function logvarexp!(out::AbstractArray, X::AbstractArray{<:Real}; corrected::Bool=true) logmeanexp!(out, X) n = _reduced_count(X, out) logsumexp!(out, _LogSqDev(X, out)) out .-= _log_count(out, max(0, corrected ? n - 1 : n)) return out end +logvarexp!(out::AbstractArray, X::AbstractArray{<:Number}; corrected::Bool=true) = _throw_not_real() -_logvarexp(X, ::Colon, corrected::Bool) = - _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) +_logvarexp(X, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) _logvarexp(X, dims, corrected::Bool) = logvarexp!(_reduced_similar(X, dims), X; corrected) """ @@ -88,40 +87,24 @@ logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = # ---- internal helpers ---- -# `log(n)` in the (real) float type of `R`, avoiding promotion (e.g. `Float32` to `Float64`). -# Works for scalar and array `R`; `n == 0` gives `-Inf`, yielding the expected `NaN`/`-Inf`. _log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) - -# number of elements reduced into each entry of `R` (empty `R` ⇒ 0, avoiding division by zero) _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) - -# output array for a reduction of `X` over `dims` (singleton along the reduced dimensions) _reduced_similar(X, dims) = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) _halve!(v::Number) = v / 2 _halve!(v::AbstractArray) = v ./= 2 -# variance/std require real inputs; reject a non-real element type up front for a clean error _throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) -_require_real(x::Real) = x -_require_real(x) = _throw_not_real() -_require_real_array(X::AbstractArray{<:Real}) = X -_require_real_array(X::AbstractArray) = _throw_not_real() -# re-iterable containers are traversed in place; any other iterator is collected once, so that -# single-use iterators (e.g. `Iterators.Stateful`) survive the variance's two passes _materialize(X) = collect(X) _materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X -# `log((exp(xᵢ) - mean)^2) = 2 * logsubexp(xᵢ, logmean)`, the term summed for the variance -_logsqdev_term(x, logmean) = 2 * logsubexp(_require_real(x), logmean) +_logsqdev_term(x::Real, logmean) = 2 * logsubexp(x, logmean) +_logsqdev_term(x, logmean) = _throw_not_real() _centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) -function _finish_logvar(logsqdev, n::Integer, corrected::Bool) - c = _log_count(logsqdev, max(0, corrected ? n - 1 : n)) - logsqdev isa Number && return logsqdev - c - logsqdev .-= c - return logsqdev +function _finish_logvar(logsqdev::Number, n::Integer, corrected::Bool) + return logsqdev - oftype(logsqdev, log(max(0, n - corrected))) end # one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` From 1fff6ba32217307c0bf42da56fdb18591598b181 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 01:07:12 +0200 Subject: [PATCH 26/40] comment --- src/logstatsexp.jl | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index c4ca9d41..8e0b5375 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -1,10 +1,5 @@ -# Numerically stable `log`-of-statistics-of-`exp` reductions. -# -# The mean is `logsumexp(X) - log(n)`. The variance uses the centered formula -# -# log var = logsumexp(2 * logsubexp(xᵢ, logmean)) - log(n - corrected) -# -# i.e. the log of the sum of squared deviations `∑ᵢ (exp(xᵢ) - mean)²`, divided by the count. +# Numerically stable `log`-of-statistics-of-`exp` reductions: +# logmeanexp, logvarexp, logstdexp. """ $(SIGNATURES) From 9983e1f3935f9f37773455407a5d89297d6ea95c Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 08:00:58 +0200 Subject: [PATCH 27/40] simplify more --- src/logstatsexp.jl | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 8e0b5375..0dfa9705 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -16,6 +16,8 @@ function logmeanexp(X) return lse - log(oftype(lse, n)) end logmeanexp(X::AbstractArray{<:Number}; dims=:) = _logmeanexp(X, dims) +_logmeanexp(X, ::Colon) = (lse = logsumexp(X); lse - log(oftype(lse, length(X)))) +_logmeanexp(X, dims) = logmeanexp!(_reduced_similar(X, dims), X) """ $(SIGNATURES) @@ -25,12 +27,9 @@ result to `out`. """ function logmeanexp!(out::AbstractArray, X::AbstractArray{<:Number}) logsumexp!(out, X) - out .-= _log_count(out, _reduced_count(X, out)) - return out + return out .-= log(_reduced_count(X, out)) end -_logmeanexp(X, ::Colon) = (lse = logsumexp(X); lse - _log_count(lse, length(X))) -_logmeanexp(X, dims) = logmeanexp!(_reduced_similar(X, dims), X) """ $(SIGNATURES) @@ -43,12 +42,14 @@ reduce over, returning `log.(var(exp.(X); dims, corrected))`. See also [`logvarexp!`](@ref). """ function logvarexp(X; corrected::Bool=true) - xs = _materialize(X) - logmean = logmeanexp(xs) - return _finish_logvar(_centered_logsqdev(xs, logmean), length(xs), corrected) + xs = collect(X) + return _finish_logvar(_centered_logsqdev(xs, logmeanexp(xs)), length(xs), corrected) end +logvarexp(X::Union{Tuple,NamedTuple,AbstractRange}; corrected::Bool=true) = _logvarexp(X, :, corrected) logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = _logvarexp(X, dims, corrected) -logvarexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = _throw_not_real() +logvarexp(X::AbstractArray; dims=:, corrected::Bool=true) = _throw_not_real() +_logvarexp(X, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) +_logvarexp(X, dims, corrected::Bool) = logvarexp!(_reduced_similar(X, dims), X; corrected) """ $(SIGNATURES) @@ -58,15 +59,12 @@ result to `out`. """ function logvarexp!(out::AbstractArray, X::AbstractArray{<:Real}; corrected::Bool=true) logmeanexp!(out, X) - n = _reduced_count(X, out) logsumexp!(out, _LogSqDev(X, out)) - out .-= _log_count(out, max(0, corrected ? n - 1 : n)) + out .-= log(max(0, _reduced_count(X, out) - corrected)) return out end logvarexp!(out::AbstractArray, X::AbstractArray{<:Number}; corrected::Bool=true) = _throw_not_real() -_logvarexp(X, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) -_logvarexp(X, dims, corrected::Bool) = logvarexp!(_reduced_similar(X, dims), X; corrected) """ $(SIGNATURES) @@ -77,12 +75,10 @@ Compute `log(std(exp, X; corrected))` in a numerically stable way. reduce over, returning `log.(std(exp.(X); dims, corrected))`. """ logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected) / 2 -logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = - _halve!(logvarexp(X; dims, corrected)) +logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = _halve!(logvarexp(X; dims, corrected)) # ---- internal helpers ---- -_log_count(R, n::Integer) = log(convert(real(float(eltype(R))), n)) _reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) _reduced_similar(X, dims) = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) @@ -91,9 +87,6 @@ _halve!(v::AbstractArray) = v ./= 2 _throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) -_materialize(X) = collect(X) -_materialize(X::Union{AbstractArray,Tuple,NamedTuple,AbstractRange}) = X - _logsqdev_term(x::Real, logmean) = 2 * logsubexp(x, logmean) _logsqdev_term(x, logmean) = _throw_not_real() _centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) From 44ddd37d09d95a4f791912c597f3eb3d53c0c12b Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 08:09:49 +0200 Subject: [PATCH 28/40] rm AbstractRange redundant --- src/logstatsexp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 0dfa9705..243d343b 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -45,7 +45,7 @@ function logvarexp(X; corrected::Bool=true) xs = collect(X) return _finish_logvar(_centered_logsqdev(xs, logmeanexp(xs)), length(xs), corrected) end -logvarexp(X::Union{Tuple,NamedTuple,AbstractRange}; corrected::Bool=true) = _logvarexp(X, :, corrected) +logvarexp(X::Union{Tuple,NamedTuple}; corrected::Bool=true) = _logvarexp(X, :, corrected) logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = _logvarexp(X, dims, corrected) logvarexp(X::AbstractArray; dims=:, corrected::Bool=true) = _throw_not_real() _logvarexp(X, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) From ef4756db5bbbc5056d4057cbd7256e50e6b2bfbd Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 08:11:41 +0200 Subject: [PATCH 29/40] _finish_logvar --- src/logstatsexp.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 243d343b..8c107daf 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -91,9 +91,7 @@ _logsqdev_term(x::Real, logmean) = 2 * logsubexp(x, logmean) _logsqdev_term(x, logmean) = _throw_not_real() _centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) -function _finish_logvar(logsqdev::Number, n::Integer, corrected::Bool) - return logsqdev - oftype(logsqdev, log(max(0, n - corrected))) -end +_finish_logvar(logsqdev::Number, n::Integer, corrected::Bool) = logsqdev - oftype(logsqdev, log(max(0, n - corrected))) # one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` function _logsumexp_count(f, X) From 172d5b855843778f02f8d04efdcce6899ddf3306 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 08:53:29 +0200 Subject: [PATCH 30/40] chainrules --- ext/LogExpFunctionsChainRulesCoreExt.jl | 63 +++++++++++++++++++++++++ test/chainrules.jl | 12 +++++ 2 files changed, 75 insertions(+) diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 28fe8fe7..35486066 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -151,6 +151,69 @@ 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=dims) + n = length(x) ÷ length(Ω) + ΔΩ = sum(exp.(x .- Ω) .* Δx; dims=dims) ./ n + return Ω, ΔΩ +end +function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims=:) + Ω = logmeanexp(x; dims=dims) + n = length(x) ÷ length(Ω) + ∂x = exp.(x .- Ω) ./ n + project_x = ChainRulesCore.ProjectTo(x) + function logmeanexp_pullback(Ω̄) + x̄ = ChainRulesCore.InplaceableThunk( + Δ -> Δ .+= Ω̄ .* ∂x, + ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), + ) + return ChainRulesCore.NoTangent(), x̄ + end + return Ω, logmeanexp_pullback +end + +function ChainRulesCore.frule((_, Δx), ::typeof(logvarexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + Ω = logvarexp(x; dims=dims, corrected=corrected) + ΔΩ = sum(_∂x_logvarexp(x, dims) .* Δx; dims=dims) + return Ω, ΔΩ +end +function ChainRulesCore.rrule(::typeof(logvarexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + Ω = logvarexp(x; dims=dims, corrected=corrected) + ∂x = _∂x_logvarexp(x, dims) + project_x = ChainRulesCore.ProjectTo(x) + function logvarexp_pullback(Ω̄) + x̄ = ChainRulesCore.InplaceableThunk( + Δ -> Δ .+= Ω̄ .* ∂x, + ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), + ) + return ChainRulesCore.NoTangent(), x̄ + end + return Ω, logvarexp_pullback +end +function ChainRulesCore.frule((_, Δx), ::typeof(logstdexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + Ω = logstdexp(x; dims=dims, corrected=corrected) + ΔΩ = sum((_∂x_logvarexp(x, dims) ./ 2) .* Δx; dims=dims) + return Ω, ΔΩ +end +function ChainRulesCore.rrule(::typeof(logstdexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true) + Ω = logstdexp(x; dims=dims, corrected=corrected) + ∂x = _∂x_logvarexp(x, dims) ./ 2 + project_x = ChainRulesCore.ProjectTo(x) + function logstdexp_pullback(Ω̄) + x̄ = ChainRulesCore.InplaceableThunk( + Δ -> Δ .+= Ω̄ .* ∂x, + ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), + ) + return ChainRulesCore.NoTangent(), x̄ + end + return Ω, logstdexp_pullback +end +function _∂x_logvarexp(x::AbstractArray{<:Real}, dims) + d = x .- logmeanexp(x; dims=dims) + e = expm1.(d) + return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims=dims) +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/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) From f1f82da6433bf2ec591cf5ce1b424393dd1cd4ac Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 12:19:42 +0200 Subject: [PATCH 31/40] no iterator support in logvarexp --- docs/src/index.md | 1 + src/LogExpFunctions.jl | 2 +- src/logstatsexp.jl | 116 ++++++++++++++++++----------------------- test/logstatsexp.jl | 63 +++++++--------------- 4 files changed, 72 insertions(+), 110 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 6c96ce3b..907b25a8 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -43,4 +43,5 @@ logmeanexp! logvarexp logvarexp! logstdexp +logstdexp! ``` diff --git a/src/LogExpFunctions.jl b/src/LogExpFunctions.jl index 19ff96a2..f9c20f09 100644 --- a/src/LogExpFunctions.jl +++ b/src/LogExpFunctions.jl @@ -9,7 +9,7 @@ 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, + logmeanexp, logmeanexp!, logvarexp, logvarexp!, logstdexp, logstdexp!, loglogistic, logitexp, log1mlogistic, logit1mexp include("basicfuns.jl") diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 8c107daf..df3a5b34 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -12,12 +12,18 @@ over, returning `log.(mean(exp.(X); dims))`. See also [`logmeanexp!`](@ref). """ function logmeanexp(X) - lse, n = _logsumexp_count(identity, X) + lse, n = _logsumexp_count(X) return lse - log(oftype(lse, n)) end logmeanexp(X::AbstractArray{<:Number}; dims=:) = _logmeanexp(X, dims) -_logmeanexp(X, ::Colon) = (lse = logsumexp(X); lse - log(oftype(lse, length(X)))) -_logmeanexp(X, dims) = logmeanexp!(_reduced_similar(X, dims), X) +function _logmeanexp(X, ::Colon) + out = logsumexp(X) + return out - oftype(out, log(length(X))) +end +function _logmeanexp(X, dims) + out = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) + return logmeanexp!(out, X) +end """ $(SIGNATURES) @@ -27,7 +33,7 @@ result to `out`. """ function logmeanexp!(out::AbstractArray, X::AbstractArray{<:Number}) logsumexp!(out, X) - return out .-= log(_reduced_count(X, out)) + return out .-= log(length(X) / length(out)) end @@ -36,20 +42,21 @@ $(SIGNATURES) Compute `log(var(exp, X; corrected))` in a numerically stable way. -`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to -reduce over, returning `log.(var(exp.(X); dims, corrected))`. +`X` should be an array of real numbers; `dims` selects the dimensions to reduce over, +returning `log.(var(exp.(X); dims, corrected))`. See also [`logvarexp!`](@ref). """ -function logvarexp(X; corrected::Bool=true) - xs = collect(X) - return _finish_logvar(_centered_logsqdev(xs, logmeanexp(xs)), length(xs), corrected) +logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims)) = _logvarexp(X, dims, corrected, logmean) +function _logvarexp(X::AbstractArray{<:Real}, dims::Colon, corrected::Bool, logmean) + out = logsumexp(2logsubexp.(X, logmean)) + return out - oftype(out, log(max(0, length(X) - corrected))) end -logvarexp(X::Union{Tuple,NamedTuple}; corrected::Bool=true) = _logvarexp(X, :, corrected) -logvarexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = _logvarexp(X, dims, corrected) -logvarexp(X::AbstractArray; dims=:, corrected::Bool=true) = _throw_not_real() -_logvarexp(X, ::Colon, corrected::Bool) = _finish_logvar(_centered_logsqdev(X, logmeanexp(X)), length(X), corrected) -_logvarexp(X, dims, corrected::Bool) = logvarexp!(_reduced_similar(X, dims), X; corrected) +function _logvarexp(X::AbstractArray{<:Real}, dims, corrected::Bool, logmean) + out = logsumexp(2logsubexp.(X, logmean); dims) + return out .-= log(max(0, length(X) / length(out) - corrected)) +end + """ $(SIGNATURES) @@ -57,72 +64,49 @@ $(SIGNATURES) Compute [`logvarexp`](@ref) of `X` over the singleton dimensions of `out`, and write the result to `out`. """ -function logvarexp!(out::AbstractArray, X::AbstractArray{<:Real}; corrected::Bool=true) +function logvarexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; corrected::Bool=true) logmeanexp!(out, X) - logsumexp!(out, _LogSqDev(X, out)) - out .-= log(max(0, _reduced_count(X, out) - corrected)) - return out + logsumexp!(out, 2logsubexp.(X, out)) + return out .-= log(max(0, length(X) / length(out) - corrected)) end -logvarexp!(out::AbstractArray, X::AbstractArray{<:Number}; corrected::Bool=true) = _throw_not_real() """ $(SIGNATURES) -Compute `log(std(exp, X; corrected))` in a numerically stable way. - -`X` should be an iterator of real numbers. For an array, `dims` selects the dimensions to -reduce over, returning `log.(std(exp.(X); dims, corrected))`. +Compute `log(std(exp, X; dims, corrected))` in a numerically stable way. """ -logstdexp(X; corrected::Bool=true) = logvarexp(X; corrected) / 2 -logstdexp(X::AbstractArray{<:Number}; dims=:, corrected::Bool=true) = _halve!(logvarexp(X; dims, corrected)) - -# ---- internal helpers ---- +logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = logvarexp(X; dims, corrected) / 2 -_reduced_count(X::AbstractArray, R) = isempty(R) ? 0 : length(X) ÷ length(R) -_reduced_similar(X, dims) = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) -_halve!(v::Number) = v / 2 -_halve!(v::AbstractArray) = v ./= 2 - -_throw_not_real() = throw(ArgumentError("logvarexp and logstdexp require real inputs")) +""" +$(SIGNATURES) -_logsqdev_term(x::Real, logmean) = 2 * logsubexp(x, logmean) -_logsqdev_term(x, logmean) = _throw_not_real() -_centered_logsqdev(X, logmean) = first(_logsumexp_count(Base.Fix2(_logsqdev_term, logmean), X)) +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 -_finish_logvar(logsqdev::Number, n::Integer, corrected::Bool) = logsqdev - oftype(logsqdev, log(max(0, n - corrected))) -# one pass over an iterator, returning `(logsumexp(f(xᵢ)), count)` -function _logsumexp_count(f, X) +# internal function to do logsumexp over an iterator in one pass, returning returning also its length +function _logsumexp_count(X) next = iterate(X) - isnothing(next) && throw(ArgumentError("reducing over an empty collection is not allowed")) - x, state = next - acc = f(x) - n = 1 - while true - next = iterate(X, state) - isnothing(next) && break + if isnothing(next) + throw(ArgumentError("reducing over an empty collection is not allowed")) + else x, state = next - acc = _logsumexp_onepass_op(acc, f(x)) - n += 1 + acc = x + n = 1 + while true + next = iterate(X, state) + isnothing(next) && break + x, state = next + acc = _logsumexp_onepass_op(acc, x) + n += 1 + end + return _logsumexp_onepass_result(acc), n end - return _logsumexp_onepass_result(acc), n -end - -# lazy array of the centered terms `2 * logsubexp(xᵢ, meanⱼ)`, with `mean` broadcast over the -# reduced (singleton) dimensions. A plain `AbstractArray` (not a `Broadcasted`) so that -# `reduce`/`reducedim!` honor offset axes and a concrete element type. -struct _LogSqDev{T,N,XT<:AbstractArray,MT<:AbstractArray} <: AbstractArray{T,N} - x::XT - mean::MT -end -_LogSqDev(x::AbstractArray, mean::AbstractArray) = - _LogSqDev{float(eltype(x)),ndims(x),typeof(x),typeof(mean)}(x, mean) -Base.size(s::_LogSqDev) = size(s.x) -Base.axes(s::_LogSqDev) = axes(s.x) -Base.IndexStyle(::Type{<:_LogSqDev}) = IndexCartesian() -Base.@propagate_inbounds function Base.getindex(s::_LogSqDev{T,N}, I::Vararg{Int,N}) where {T,N} - j = map((i, ax) -> ifelse(length(ax) == 1, first(ax), i), I, axes(s.mean)) - return convert(T, _logsqdev_term(s.x[I...], s.mean[j...])) end diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 4722bd0b..5306e4d1 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -1,7 +1,7 @@ using Test: @test, @test_throws, @testset, @inferred using Statistics: mean, std, var using OffsetArrays: OffsetArray -using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logvarexp, logvarexp! +using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logstdexp!, logvarexp, logvarexp! # Count heap allocations of `f(x)` after warming it up. `f` and `x` are passed as # arguments so that they are concretely typed inside this function (avoiding spurious @@ -43,35 +43,22 @@ end end end -@testset "logmeanexp, logvarexp, logstdexp iterators" begin +@testset "logmeanexp iterators, empty reductions" begin x = randn(Float32, 20) xt = Tuple(x) xg = (v for v in x) xf = Iterators.filter(_ -> true, x) - xe = exp.(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)) - @test @inferred(logvarexp(xt)) ≈ log(var(xe)) - @test logvarexp(xt; corrected=false) ≈ log(var(xe; corrected=false)) - @test @inferred(logvarexp((v for v in x))) ≈ log(var(xe)) - @test logvarexp((v for v in x); corrected=false) ≈ log(var(xe; corrected=false)) - @test logvarexp(Iterators.Stateful(x)) ≈ log(var(xe)) - @test @inferred(logstdexp(xt)) ≈ log(std(xe)) - @test logstdexp(xt; corrected=false) ≈ log(std(xe; corrected=false)) - @test @inferred(logstdexp((v for v in x))) ≈ log(std(xe)) - @test logstdexp((v for v in x); corrected=false) ≈ log(std(xe; corrected=false)) - @test logstdexp(Iterators.Stateful(x)) ≈ log(std(xe)) - @test isnan(logvarexp((0.0,))) - @test isnan(logstdexp((0.0,))) - @test_throws ArgumentError logvarexp((1.0 + 0.0im, 2.0 + 0.0im)) - @test_throws ArgumentError logstdexp((1.0 + 0.0im, 2.0 + 0.0im)) + + # empty `Tuple` has unknown eltype, so the reduction errors; empty arrays give NaN @test_throws ArgumentError logmeanexp(()) - @test isnan(logmeanexp(Float64[])) # empty array mean is NaN, matching `mean` - @test_throws ArgumentError logvarexp(()) - @test_throws ArgumentError logvarexp(Float64[]) + @test isnan(logmeanexp(Float64[])) + @test isnan(logvarexp(Float64[])) end @testset "logmeanexp, logvarexp, logstdexp promotion and dims coverage" begin @@ -140,12 +127,12 @@ end data = randn(7) @test logmeanexp(DrainOnce(data)) ≈ log(mean(exp, data)) - # Complex arrays are rejected with a clear ArgumentError on every variance/std path, - # with or without `dims` (previously the `dims` form threw a confusing MethodError). + # Complex arrays are rejected (no `<:Real` method) on every variance/std path, with or + # without `dims`; some error is thrown (the exact exception type is unspecified). C = ComplexF64[1 2; 3 4] - @test_throws ArgumentError logvarexp(C) - @test_throws ArgumentError logvarexp(C; dims=1) - @test_throws ArgumentError logstdexp(C; dims=2) + @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!" begin @@ -164,29 +151,22 @@ 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_throws ArgumentError logvarexp!(Matrix{ComplexF64}(undef, 1, 2), ComplexF64[1 2; 3 4]) + @test_throws Exception logvarexp!(Matrix{ComplexF64}(undef, 1, 2), ComplexF64[1 2; 3 4]) end @testset "allocations" begin - # Full reductions allocate at most a small constant (no per-element / O(n) temporary): - # the allocation count must not grow with the input size. + # `logmeanexp`'s `dims=:`, `dims`, and in-place paths build no O(n) temporary, so + # allocations don't grow with input size. (`logvarexp`/`logstdexp` do allocate an O(n) + # `2logsubexp.(X, logmean)` temporary — a deliberate simplicity tradeoff — so they + # aren't checked here.) for T in (Float32, Float64) - for f in (logmeanexp, logvarexp, logstdexp) - @test allocations(f, randn(T, 10_000)) == allocations(f, randn(T, 100)) - # `dims` reductions only allocate the (output-sized) result and scratch — never an - # O(n) temporary — so allocations are independent of the reduced dimension's length. - g = X -> f(X; dims=1) - @test allocations(g, randn(T, 10_000, 3)) == allocations(g, randn(T, 100, 3)) - end - # genuinely allocation-free paths + @test allocations(logmeanexp, randn(T, 10_000)) == allocations(logmeanexp, randn(T, 100)) @test allocations(logmeanexp, randn(T, 10_000)) == 0 - @test allocations(logvarexp, Tuple(randn(T, 20))) == 0 - # in-place reductions reuse `out`, so their allocation does not grow with the input + 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)) - @test allocations(X -> logvarexp!(out, X), randn(T, 10_000, 3)) == - allocations(X -> logvarexp!(out, X), randn(T, 100, 3)) end end @@ -215,8 +195,5 @@ end @test logstdexp(x; corrected) ≈ refvar(x; corrected) / 2 rtol = 1e-8 atol = 1e-9 end end - # array and one-shot-iterator paths agree even in the cancellation regime - x = 1.0 .+ 1e-5 .* randn(100) - @test logvarexp(x) ≈ logvarexp(Iterators.Stateful(x)) rtol = 1e-10 end end From 0a900f08dea3d7edf7e821bfbd17a2afc14d9015 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 12:20:47 +0200 Subject: [PATCH 32/40] logstdexp! tests --- test/logstatsexp.jl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 5306e4d1..731824be 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -135,7 +135,7 @@ end @test_throws Exception logstdexp(C; dims=2) end -@testset "in-place logmeanexp!/logvarexp!" begin +@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)) @@ -146,12 +146,17 @@ end 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 From 850df334b879ae1933acbbdccef6fb39b9941dd6 Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 17:17:11 +0200 Subject: [PATCH 33/40] nits --- src/logstatsexp.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index df3a5b34..4895f1ce 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -66,7 +66,7 @@ result to `out`. """ function logvarexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; corrected::Bool=true) logmeanexp!(out, X) - logsumexp!(out, 2logsubexp.(X, out)) + logsumexp!(out, 2 .* logsubexp.(X, out)) return out .-= log(max(0, length(X) / length(out) - corrected)) end @@ -76,7 +76,7 @@ $(SIGNATURES) Compute `log(std(exp, X; dims, corrected))` in a numerically stable way. """ -logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true) = logvarexp(X; dims, corrected) / 2 +logstdexp(X::AbstractArray{<:Real}; dims=:, corrected::Bool=true, logmean=logmeanexp(X; dims)) = logvarexp(X; dims, corrected, logmean) / 2 """ @@ -91,7 +91,7 @@ function logstdexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; correc end -# internal function to do logsumexp over an iterator in one pass, returning returning also its length +# internal function to do logsumexp over an iterator in one pass, returning also its length function _logsumexp_count(X) next = iterate(X) if isnothing(next) From b245d7df15e9119c194b75873e7979bccb58009e Mon Sep 17 00:00:00 2001 From: cossio Date: Tue, 9 Jun 2026 20:06:22 +0200 Subject: [PATCH 34/40] simplify pullback --- ext/LogExpFunctionsChainRulesCoreExt.jl | 62 ++++++++++--------------- src/logstatsexp.jl | 36 +++++++------- 2 files changed, 43 insertions(+), 55 deletions(-) diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 35486066..90344527 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -152,66 +152,54 @@ function ChainRulesCore.rrule(::typeof(logsumexp), x::AbstractArray{<:Real}; dim end function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims=:) - Ω = logmeanexp(x; dims=dims) + Ω = logmeanexp(x; dims) n = length(x) ÷ length(Ω) - ΔΩ = sum(exp.(x .- Ω) .* Δx; dims=dims) ./ n + ΔΩ = sum(exp.(x .- Ω) .* Δx; dims) ./ n return Ω, ΔΩ end function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims=:) - Ω = logmeanexp(x; dims=dims) + Ω = logmeanexp(x; dims) n = length(x) ÷ length(Ω) - ∂x = exp.(x .- Ω) ./ n - project_x = ChainRulesCore.ProjectTo(x) - function logmeanexp_pullback(Ω̄) - x̄ = ChainRulesCore.InplaceableThunk( - Δ -> Δ .+= Ω̄ .* ∂x, - ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), - ) - return ChainRulesCore.NoTangent(), x̄ - end - return Ω, logmeanexp_pullback + return Ω, _∂x_pullback(exp.(x .- Ω) ./ n, x) end function ChainRulesCore.frule((_, Δx), ::typeof(logvarexp), x::AbstractArray{<:Real}; dims=:, corrected::Bool=true) - Ω = logvarexp(x; dims=dims, corrected=corrected) - ΔΩ = sum(_∂x_logvarexp(x, dims) .* Δx; dims=dims) + 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) - Ω = logvarexp(x; dims=dims, corrected=corrected) - ∂x = _∂x_logvarexp(x, dims) - project_x = ChainRulesCore.ProjectTo(x) - function logvarexp_pullback(Ω̄) - x̄ = ChainRulesCore.InplaceableThunk( - Δ -> Δ .+= Ω̄ .* ∂x, - ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), - ) - return ChainRulesCore.NoTangent(), x̄ - end - return Ω, logvarexp_pullback + 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) - Ω = logstdexp(x; dims=dims, corrected=corrected) - ΔΩ = sum((_∂x_logvarexp(x, dims) ./ 2) .* Δx; dims=dims) + 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) - Ω = logstdexp(x; dims=dims, corrected=corrected) - ∂x = _∂x_logvarexp(x, dims) ./ 2 + 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 logstdexp_pullback(Ω̄) + function pullback(Ω̄) x̄ = ChainRulesCore.InplaceableThunk( Δ -> Δ .+= Ω̄ .* ∂x, ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), ) return ChainRulesCore.NoTangent(), x̄ end - return Ω, logstdexp_pullback -end -function _∂x_logvarexp(x::AbstractArray{<:Real}, dims) - d = x .- logmeanexp(x; dims=dims) - e = expm1.(d) - return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims=dims) + return pullback end # no rules for mutating functions currently: diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 4895f1ce..1dde3e63 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -40,20 +40,21 @@ end """ $(SIGNATURES) -Compute `log(var(exp, X; corrected))` in a numerically stable way. +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))`. +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}, dims::Colon, corrected::Bool, logmean) - out = logsumexp(2logsubexp.(X, logmean)) + out = logsumexp(2 .* logsubexp.(X, logmean)) return out - oftype(out, log(max(0, length(X) - corrected))) end function _logvarexp(X::AbstractArray{<:Real}, dims, corrected::Bool, logmean) - out = logsumexp(2logsubexp.(X, logmean); dims) + out = logsumexp(2 .* logsubexp.(X, logmean); dims) return out .-= log(max(0, length(X) / length(out) - corrected)) end @@ -74,7 +75,9 @@ end """ $(SIGNATURES) -Compute `log(std(exp, X; dims, corrected))` in a numerically stable way. +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 @@ -94,19 +97,16 @@ end # internal function to do logsumexp over an iterator in one pass, returning also its length function _logsumexp_count(X) next = iterate(X) - if isnothing(next) - throw(ArgumentError("reducing over an empty collection is not allowed")) - else + isnothing(next) && throw(ArgumentError("reducing over an empty collection is not allowed")) + x, state = next + acc = x + n = 1 + while true + next = iterate(X, state) + isnothing(next) && break x, state = next - acc = x - n = 1 - while true - next = iterate(X, state) - isnothing(next) && break - x, state = next - acc = _logsumexp_onepass_op(acc, x) - n += 1 - end - return _logsumexp_onepass_result(acc), n + acc = _logsumexp_onepass_op(acc, x) + n += 1 end + return _logsumexp_onepass_result(acc), n end From 0ce6af474299afeda711420a7010b9b38ea965ee Mon Sep 17 00:00:00 2001 From: cossio Date: Wed, 10 Jun 2026 14:56:28 +0200 Subject: [PATCH 35/40] rm _logsumexp_count --- src/logstatsexp.jl | 40 +++++++++++++++++----------------------- 1 file changed, 17 insertions(+), 23 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index 1dde3e63..a92054c4 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -12,15 +12,27 @@ over, returning `log.(mean(exp.(X); dims))`. See also [`logmeanexp!`](@ref). """ function logmeanexp(X) - lse, n = _logsumexp_count(X) - return lse - log(oftype(lse, n)) + 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 lse - log(oftype(lse, count)) end logmeanexp(X::AbstractArray{<:Number}; dims=:) = _logmeanexp(X, dims) -function _logmeanexp(X, ::Colon) +function _logmeanexp(X::AbstractArray{<:Number}, ::Colon) out = logsumexp(X) return out - oftype(out, log(length(X))) end -function _logmeanexp(X, dims) +function _logmeanexp(X::AbstractArray{<:Number}, dims) out = similar(X, float(eltype(X)), Base.reduced_indices(axes(X), dims)) return logmeanexp!(out, X) end @@ -49,7 +61,7 @@ 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}, dims::Colon, corrected::Bool, logmean) +function _logvarexp(X::AbstractArray{<:Real}, ::Colon, corrected::Bool, logmean) out = logsumexp(2 .* logsubexp.(X, logmean)) return out - oftype(out, log(max(0, length(X) - corrected))) end @@ -92,21 +104,3 @@ function logstdexp!(out::AbstractArray{<:Real}, X::AbstractArray{<:Real}; correc logvarexp!(out, X; corrected) return out ./= 2 end - - -# internal function to do logsumexp over an iterator in one pass, returning also its length -function _logsumexp_count(X) - next = iterate(X) - isnothing(next) && throw(ArgumentError("reducing over an empty collection is not allowed")) - x, state = next - acc = x - n = 1 - while true - next = iterate(X, state) - isnothing(next) && break - x, state = next - acc = _logsumexp_onepass_op(acc, x) - n += 1 - end - return _logsumexp_onepass_result(acc), n -end From e6cc62a9ac2026b4c429337d39b6e79a1a8d148b Mon Sep 17 00:00:00 2001 From: cossio Date: Sun, 21 Jun 2026 14:13:08 +0200 Subject: [PATCH 36/40] Simplify oftype usage and add return-type tests Hoist `oftype` to wrap the full expression in the scalar logmeanexp / logvarexp reductions, and add a `return types` testset asserting Float32->Float32, Float64->Float64, and Int->Float64 across the scalar, dims, dims=:, iterator, and in-place paths (via @inferred). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1 --- src/logstatsexp.jl | 6 +++--- test/logstatsexp.jl | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/src/logstatsexp.jl b/src/logstatsexp.jl index a92054c4..f8fca575 100644 --- a/src/logstatsexp.jl +++ b/src/logstatsexp.jl @@ -25,12 +25,12 @@ function logmeanexp(X) count += 1 end lse = _logsumexp_onepass_result(acc) - return lse - log(oftype(lse, count)) + 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 out - oftype(out, log(length(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)) @@ -63,7 +63,7 @@ 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 out - oftype(out, log(max(0, length(X) - corrected))) + 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) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 731824be..ea418b8f 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -82,6 +82,42 @@ end @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 return a 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 return an array whose eltype is `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 (the edited `oftype(lse, lse - log(count))` line) + @test typeof(@inferred(logmeanexp(Tuple(vec(X))))) == Tout + end + + # in-place reductions write into the caller's `out`, so they preserve its eltype — + # including the widening case where an integer input is reduced 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 + # Regressions for correctness bugs found in review. Each block is one bug class. @testset "edge-case regressions" begin # Non-1-based axes (OffsetArrays): the `dims` variance/std must match the result on From 82d59b1d4938477b429cec0eb8725bb660f88491 Mon Sep 17 00:00:00 2001 From: cossio Date: Sun, 21 Jun 2026 14:13:36 +0200 Subject: [PATCH 37/40] comments --- .claude/worktrees/logvarexp-alloc | 1 + .claude/worktrees/warm-dazzling-avalanche | 1 + test/logstatsexp.jl | 6 ++---- 3 files changed, 4 insertions(+), 4 deletions(-) create mode 160000 .claude/worktrees/logvarexp-alloc create mode 160000 .claude/worktrees/warm-dazzling-avalanche diff --git a/.claude/worktrees/logvarexp-alloc b/.claude/worktrees/logvarexp-alloc new file mode 160000 index 00000000..0ce6af47 --- /dev/null +++ b/.claude/worktrees/logvarexp-alloc @@ -0,0 +1 @@ +Subproject commit 0ce6af474299afeda711420a7010b9b38ea965ee diff --git a/.claude/worktrees/warm-dazzling-avalanche b/.claude/worktrees/warm-dazzling-avalanche new file mode 160000 index 00000000..f7fd06af --- /dev/null +++ b/.claude/worktrees/warm-dazzling-avalanche @@ -0,0 +1 @@ +Subproject commit f7fd06af7cdb1824bb459cdd3e78bbdb0c29d23e diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index ea418b8f..9abc949e 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -3,14 +3,12 @@ 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. `f` and `x` are passed as -# arguments so that they are concretely typed inside this function (avoiding spurious -# allocations from captured, boxed locals). +# Count heap allocations of `f(x)` after warming it up allocations(f, x) = (f(x); @allocated f(x)) # A single-use iterator that nonetheless advertises a length (so it exercises the # length-aware code path). Iterating it consumes it; a second traversal — or an -# `isempty` probe — would skip elements. Used to check `logmeanexp`'s one-pass contract. +# `isempty` probe — would skip elements. Used to check `logmeanexp`'s one-pass. mutable struct DrainOnce{T} data::Vector{T} pos::Int From 69f455ba8090155e3c4cd917c07cccdb7a8afb76 Mon Sep 17 00:00:00 2001 From: cossio Date: Sun, 21 Jun 2026 14:20:09 +0200 Subject: [PATCH 38/40] Shorten test comments Tighten the multi-line explanatory comments in test/logstatsexp.jl while keeping the rationale each one records. Comment-only; no test logic changed. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1 --- test/logstatsexp.jl | 49 ++++++++++++++++++--------------------------- 1 file changed, 19 insertions(+), 30 deletions(-) diff --git a/test/logstatsexp.jl b/test/logstatsexp.jl index 9abc949e..ec2b8dad 100644 --- a/test/logstatsexp.jl +++ b/test/logstatsexp.jl @@ -6,9 +6,8 @@ using LogExpFunctions: logmeanexp, logmeanexp!, logstdexp, logstdexp!, logvarexp # Count heap allocations of `f(x)` after warming it up allocations(f, x) = (f(x); @allocated f(x)) -# A single-use iterator that nonetheless advertises a length (so it exercises the -# length-aware code path). Iterating it consumes it; a second traversal — or an -# `isempty` probe — would skip elements. Used to check `logmeanexp`'s one-pass. +# 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 @@ -85,7 +84,7 @@ end 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 return a scalar of type `Tout` + # 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 @@ -93,19 +92,18 @@ end @test typeof(@inferred(logstdexp(X))) == Tout @test typeof(@inferred(logstdexp(X; dims=:))) == Tout - # `dims` reductions return an array whose eltype is `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 (the edited `oftype(lse, lse - log(count))` line) + # single-pass iterator path @test typeof(@inferred(logmeanexp(Tuple(vec(X))))) == Tout end - # in-place reductions write into the caller's `out`, so they preserve its eltype — - # including the widening case where an integer input is reduced into a float `out`. + # 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 @@ -116,10 +114,10 @@ end @test eltype(logstdexp!(Matrix{Float32}(undef, 1, 4), Xf)) == Float32 end -# Regressions for correctness bugs found in review. Each block is one bug class. +# Regression tests for review-found bugs, one block per bug class. @testset "edge-case regressions" begin - # Non-1-based axes (OffsetArrays): the `dims` variance/std must match the result on - # the equivalent 1-based array (a fused lazy reduction silently returned wrong values). + # 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, :) @@ -132,16 +130,14 @@ end end end - # Abstract element type: the `dims` variance must still work (and match a concretely - # typed copy), not throw a MethodError. + # 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 reduction along the reduced dimension: `var` is NaN (not an error) for every - # `corrected`, matching `Statistics.var` and `logmeanexp`. + # 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) @@ -149,20 +145,17 @@ end @test all(isnan, logstdexp(Eredux; dims=1, corrected)) end - # Empty along a dimension that is NOT being reduced: the result is an empty array of - # the reduced shape (no DivideError). + # 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 that reports a length: logmeanexp must traverse it exactly once - # (a length-aware path that re-consumed it would skip the first element). + # 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, with or - # without `dims`; some error is thrown (the exact exception type is unspecified). + # 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) @@ -194,10 +187,8 @@ end end @testset "allocations" begin - # `logmeanexp`'s `dims=:`, `dims`, and in-place paths build no O(n) temporary, so - # allocations don't grow with input size. (`logvarexp`/`logstdexp` do allocate an O(n) - # `2logsubexp.(X, logmean)` temporary — a deliberate simplicity tradeoff — so they - # aren't checked here.) + # `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 @@ -210,9 +201,8 @@ end end @testset "numerical robustness" begin - # Compare against high-precision (BigFloat) references on hard cases: tight clusters - # (var ≪ mean², where a raw second-moment formula cancels catastrophically) and - # values large/small enough that exp would over-/under-flow Float64. + # 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))) @@ -226,8 +216,7 @@ end [0.0, 1e-6], ) for x in cases - # atol covers near-zero results (e.g. logmeanexp([0, 1e-6]) ≈ 5e-7), where - # an inherent cancellation makes the relative error meaningless + # 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 From 310aecb2fbe00fb8e49dcd80b4f60895df7915ea Mon Sep 17 00:00:00 2001 From: cossio Date: Sun, 21 Jun 2026 14:20:09 +0200 Subject: [PATCH 39/40] Stop tracking .claude worktrees The previous commit accidentally added two local Claude Code worktree directories as gitlinks (mode 160000); remove them and ignore .claude/. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1 --- .claude/worktrees/logvarexp-alloc | 1 - .claude/worktrees/warm-dazzling-avalanche | 1 - .gitignore | 1 + 3 files changed, 1 insertion(+), 2 deletions(-) delete mode 160000 .claude/worktrees/logvarexp-alloc delete mode 160000 .claude/worktrees/warm-dazzling-avalanche diff --git a/.claude/worktrees/logvarexp-alloc b/.claude/worktrees/logvarexp-alloc deleted file mode 160000 index 0ce6af47..00000000 --- a/.claude/worktrees/logvarexp-alloc +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 0ce6af474299afeda711420a7010b9b38ea965ee diff --git a/.claude/worktrees/warm-dazzling-avalanche b/.claude/worktrees/warm-dazzling-avalanche deleted file mode 160000 index f7fd06af..00000000 --- a/.claude/worktrees/warm-dazzling-avalanche +++ /dev/null @@ -1 +0,0 @@ -Subproject commit f7fd06af7cdb1824bb459cdd3e78bbdb0c29d23e diff --git a/.gitignore b/.gitignore index 045f79b8..1c862d24 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ Manifest.toml /deps/deps.jl /docs/build +.claude/ From 899b08f469bb30c498069b48406b9f5445816b46 Mon Sep 17 00:00:00 2001 From: cossio Date: Sun, 21 Jun 2026 17:05:57 +0200 Subject: [PATCH 40/40] Remove .claude from .gitignore Keep .claude out of the tracked .gitignore; it is ignored locally via .git/info/exclude instead (per-clone, not committed). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1 --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 1c862d24..045f79b8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,3 @@ Manifest.toml /deps/deps.jl /docs/build -.claude/