Add ChainRulesCore extension with frule/rrule for logmeanexp, logvarexp, logstdexp - #3
Conversation
Define forward- and reverse-mode ChainRules for logmeanexp, logvarexp and logstdexp in a ChainRulesCore package extension, following the pattern of JuliaStats/LogExpFunctions.jl#120. The logvarexp/logstdexp rules accept the same logmean keyword as the primal functions. Tested against finite differences with ChainRulesTestUtils across dims and corrected combinations. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 1 2 +1
Lines 24 59 +35
=========================================
+ Hits 24 59 +35 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f7e85d7f04
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| @testset "chainrules" begin | ||
| 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 |
There was a problem hiding this comment.
When the loop reaches dims = 1 or dims = 2, this condition attempts to iterate an Int, which raises a MethodError in Julia before any derivative checks run. Normalize an integer dimension to a one-element tuple, or handle dims isa Integer separately, so the new test suite can complete.
Useful? React with 👍 / 👎.
| d = x .- logmean | ||
| e = expm1.(d) | ||
| return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims) |
There was a problem hiding this comment.
For nearly equal but representably distinct inputs, squaring e can underflow even though both the primal and its derivative are finite; for example, with Float64[-1e-200, 1e-200], sum(abs2, e) becomes zero while the expected log-variance gradient is on the order of 1e200, so these rules return Inf/NaN. Compute the ratio with log-domain arithmetic or scale e before forming the sum of squares.
Useful? React with 👍 / 👎.
| Δ -> Δ .+= Ω̄ .* ∂x, | ||
| ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), |
There was a problem hiding this comment.
When an upstream ChainRules pullback supplies Ω̄ as a supported Thunk, both branches capture that wrapper and later attempt Ω̄ .* ∂x; evaluating the returned tangent then fails because the thunk itself is not the cotangent array or scalar. Unthunk Ω̄ before using it, or explicitly compose the thunks so this rule works in reverse-mode chains that preserve lazy cotangents.
Useful? React with 👍 / 👎.
| using LogStatFunctions: logmeanexp, logvarexp, logstdexp | ||
| import ChainRulesCore | ||
|
|
||
| function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :) |
There was a problem hiding this comment.
The package already supports complex arrays for logmeanexp, but every new rule is restricted to AbstractArray{<:Real} while the added documentation says all three functions have ChainRules derivatives without stating that limitation. Consequently, rrule(logmeanexp, complex_x) has no method from this extension and complex users do not receive the advertised AD support; either add the correctly conjugated complex rule or document the real-input restriction.
Useful? React with 👍 / 👎.
Adds ChainRules support via a
ChainRulesCorepackage extension, following the pattern of JuliaStats/LogExpFunctions.jl#120.Changes
ext/LogStatFunctionsChainRulesCoreExt.jldefiningfrules andrrules forlogmeanexp,logvarexpandlogstdexp(withdimsandcorrectedkeyword support). Thelogvarexp/logstdexprules also accept thelogmeankeyword, matching the primal signatures.logvarexpgradient uses∂/∂xⱼ = 2 exp(xⱼ)(exp(xⱼ) − m) / Σᵢ(exp(xᵢ) − m)²withm = exp(logmean); the dependence ofmonxcancels becauseΣᵢ(exp(xᵢ) − m) = 0, andcorrectedonly shifts the result by a constant.logstdexpis half of that. Reverse rules return anInplaceableThunkwithProjectTo, as in the LogExpFunctions PR.Project.toml:ChainRulesCoreas a weak dependency with compat"1".test/chainrules.jlchecksfrule/rruleagainst finite differences with ChainRulesTestUtils over vectors/matrices,dims ∈ (:, 1, 1:2, 2)and bothcorrectedvalues (330 assertions).Verification
Pkg.test()passes locally on Julia 1.12.6 (functions, chainrules, ExplicitImports, Aqua).--checkpasses.🤖 Generated with Claude Code
https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj