logstatsexp v2 by cossio · Pull Request #120 · JuliaStats/LogExpFunctions.jl

cossio · 2026-06-03T20:24:56Z

Adds the functions logmeanexp, logvarexp, logstdexp

A second attempt at #77 . The content is the same, but with some cleanup addressing comments at #77 (with agent help). @tpapp

Fixes #81

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds numerically-stable “log-statistics of exp” helpers (mean/variance/std) to LogExpFunctions, along with tests and docs coverage.

Changes:

Export and include new APIs: logmeanexp, logvarexp, logstdexp
Add test coverage for arrays/iterators, dims, corrected, and type inference/promotion
Document the new functions in the docs index

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
test/runtests.jl	Registers the new `logstatsexp` test file in the test suite.
test/logstatsexp.jl	Adds correctness and inference tests for the new log-statistic functions.
test/Project.toml	Adds `Statistics` as a test dependency to support reference computations.
src/logstatsexp.jl	Implements the new `logmeanexp` / `logvarexp` / `logstdexp` APIs and internal helpers.
src/LogExpFunctions.jl	Exports and includes the new implementation file.
docs/src/index.md	Lists the newly exported functions for documentation.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…ocstrings to say real numbers

tpapp

Thanks @cossio for coming back to this.

I find the copilot-generated code a bit convoluted. Specifically,

I am not sure why it makes sense to accumulate the length in general, when array-like objects can compute it cheaply; the only context where that makes sense is a generic iterator with Base.SizeUnknown().
I am not sure of the purpose of _subtract_log_count, shouldn't -Inf - log(0) just take care of the empty case?

Generally, the way I think could be coded is by defining an

initial state,
an elementwise op,
a finalizer that gives the results
for each of these.

Then an internal function could implement something like

function _reduce(op, a::AbstractArray{T}, dims)
    _finalize(reduce(_reduction(op), a; dims, init = _init(op, float(T)), a)
end

where _finalize(::typeof(logmeanexp)) etc are defined. A trait mechanism could special-case the Base.SizeUnknown() iterator.

Note that this is just a suggestion, you should code it the way you prefer. I just want to keep the implementation simple and easy to maintain.

Incidentally, mean_and_std and mean_and_var would make sense.

@tpapp

…ogstdexp Rework the log-mean/var/std-exp reductions around a single primitive (`logsumexp(X)`, `logsumexp(2X)`, count), addressing review feedback in JuliaStats#120 and JuliaStats#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) <noreply@anthropic.com>

cossio · 2026-06-04T15:59:26Z

Added logmeanexp_and_logvarexp and logmeanexp_and_logstdexp, which share a common pass and thus are roughly 1.5x faster than individual calls to logmeanexp and logvarexp.

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) <noreply@anthropic.com>

`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) <noreply@anthropic.com>

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

`_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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

cossio · 2026-06-08T16:12:29Z

@tpapp friendly bump, anything else needed here ?

devmotion

IMO this is a terribly verbose and partly overly complicated PR - definitely subpar to any existing implementations in this package. I think this should be cleaned up quite a bit before being merged.

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

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) <noreply@anthropic.com>

…e output 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) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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) <noreply@anthropic.com>

cossio · 2026-06-09T06:14:35Z

@devmotion cleaned-up a bit more. Hopefully it's better now.

Also added AD rules.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

devmotion · 2026-06-09T08:05:26Z

+# 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


Is this really necessary? Since Julia doesn't have any contracts for (arguably loosely defined) interfaces, it's just too easy to miss definitions or define things in inefficient ways. I've seen this happen in particular with custom array definitions quite a bit...

I somehow have a slightly hopeless feeling regarding this PR - it seems to overcomplicate things and contain code that I (maybe incorrectly) assume was introduced by an AI system without controlling or thinking through the design. That gives me bad vibes when reviewing and will be annoying and unnecessarily complex for future maintainers.

@devmotion This _LogSqDev construct was there to try to reduce allocations in logvarexp!. This was suggested by Claude, but I myself see no alternative way to reduce allocations logvarexp!.

If we are willing to accept a temporary allocation in logvarexp! then I can remove _LogSqDev it and I agree the code will be cleaner.

If you agree, I'll also drop support for iterators in logvarexp. The complexity is because it needs two passes over the iterator, so one needs to materialize the iterator anyways. I think with these two modifications will simplify the code substantially.

We then support iterators only for logmeanexp which can be done in one-pass. As far as I can see this is the cleanest solution.

I implemented the changes above:

removed _LogSqDev

dropped support for iterators in logvarexp

Could you take another look ? I think it's much cleaner now. Thanks

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) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1

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) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1

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) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1

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) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0161bYmHCny8nuA4agW2KnJ1

Copilot AI added 4 commits June 3, 2026 18:48

Add logmeanexp/logvarexp/logstdexp implementations

5b96b87

Fix logstats iterator implementation and tests

aa991ca

Add PR77-focused logstatsexp tests

3930a78

Fix PR review issues in logstatsexp

0431c84

Copilot AI review requested due to automatic review settings June 3, 2026 20:24

cossio mentioned this pull request Jun 3, 2026

logstatsexp #77

Closed

Copilot AI reviewed Jun 3, 2026

View reviewed changes

Comment thread src/logstatsexp.jl Outdated

Comment thread src/logstatsexp.jl Outdated

Copilot AI added 3 commits June 3, 2026 20:39

Fix unresolved review comments: remove dead _count_elements, update d…

6c42836

…ocstrings to say real numbers

Use explicit multiplication in logstatsexp

7f14412

fix iterator real-input validation for log variance

28fc64e

tpapp requested changes Jun 4, 2026

View reviewed changes

Copilot AI mentioned this pull request Jun 4, 2026

Simplify logstatsexp reductions: length-aware iterators and direct normalization cossio/LogExpFunctions.jl#1

Open

Copilot AI and others added 2 commits June 4, 2026 11:57

refactor iterator counting and simplify normalization logic

7a886c1

cossio and others added 7 commits June 4, 2026 18:31

Potential fix for pull request finding

8c421b5

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

73f2e55

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

cossio requested a review from tpapp June 4, 2026 21:33

tpapp reviewed Jun 5, 2026

View reviewed changes

Comment thread src/logstatsexp.jl Outdated

devmotion reviewed Jun 8, 2026

View reviewed changes

cossio and others added 3 commits June 8, 2026 19:56

Apply suggestion from @devmotion

8fafe93

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

Apply suggestion from @devmotion

c9439aa

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

Apply suggestion from @devmotion

1be77c1

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

cossio force-pushed the copilot/retry-pr-77 branch from 01d92c8 to 691e34e Compare June 8, 2026 18:15

cossio and others added 8 commits June 8, 2026 21:47

Use keyword-argument punning in logstatsexp tests

5601c23

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

simplifications

8de5fea

comment

1fff6ba

simplify more

9983e1f

rm AbstractRange redundant

44ddd37

_finish_logvar

ef4756d

chainrules

172d5b8

cossio requested a review from Copilot June 9, 2026 06:54

Copilot AI reviewed Jun 9, 2026

View reviewed changes

Comment thread ext/LogExpFunctionsChainRulesCoreExt.jl Outdated

Comment thread ext/LogExpFunctionsChainRulesCoreExt.jl

Comment thread ext/LogExpFunctionsChainRulesCoreExt.jl Outdated

Comment thread src/logstatsexp.jl Outdated

devmotion reviewed Jun 9, 2026

View reviewed changes

cossio and others added 10 commits June 9, 2026 12:19

no iterator support in logvarexp

f1f82da

logstdexp! tests

0a900f0

nits

850df33

simplify pullback

b245d7d

rm _logsumexp_count

0ce6af4

comments

82d59b1

Uh oh!

Conversation

cossio commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

tpapp left a comment

Choose a reason for hiding this comment

Uh oh!

cossio commented Jun 4, 2026

Uh oh!

Uh oh!

cossio commented Jun 8, 2026

Uh oh!

devmotion left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cossio commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

devmotion Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

cossio Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cossio Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cossio Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

cossio commented Jun 3, 2026 •

edited

Loading

cossio commented Jun 9, 2026 •

edited

Loading

cossio Jun 9, 2026 •

edited

Loading

cossio Jun 9, 2026 •

edited

Loading

cossio Jun 9, 2026 •

edited

Loading