From f7e85d7f04b8d7fe6ab0244c867d94eba6b8d450 Mon Sep 17 00:00:00 2001 From: cossio Date: Wed, 29 Jul 2026 12:18:19 +0200 Subject: [PATCH] Add ChainRulesCore extension with frule/rrule for all functions 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 Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj --- CHANGELOG.md | 3 + Project.toml | 7 +++ docs/src/index.md | 5 ++ ext/LogStatFunctionsChainRulesCoreExt.jl | 74 ++++++++++++++++++++++++ test/Project.toml | 2 + test/chainrules.jl | 17 ++++++ test/runtests.jl | 3 + 7 files changed, 111 insertions(+) create mode 100644 ext/LogStatFunctionsChainRulesCoreExt.jl create mode 100644 test/chainrules.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a6826e..16650fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ All notable changes to this project will be documented in this file. The format ## Unreleased +- Add a ChainRulesCore extension defining `frule`s and `rrule`s for + `logmeanexp`, `logvarexp` and `logstdexp`, making them differentiable with + ChainRules-based AD packages (e.g. Zygote). - Support immutable arrays (e.g. `StaticArrays`) and complex arrays. - Results preserve the input eltype (e.g. `Float32` in → `Float32` out), and the `log(N)` normalization is computed in the result's precision (so e.g. diff --git a/Project.toml b/Project.toml index 9ce0ac2..37bc4bf 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,13 @@ projects = ["test", "docs"] [deps] LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +[weakdeps] +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + +[extensions] +LogStatFunctionsChainRulesCoreExt = "ChainRulesCore" + [compat] +ChainRulesCore = "1" LogExpFunctions = "0.3.26, 1.0" julia = "1.10" diff --git a/docs/src/index.md b/docs/src/index.md index 482a7ee..874a37d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -28,6 +28,11 @@ logstdexp(A; corrected=false) `logvarexp` and `logstdexp` also accept `logmean` to reuse a precomputed `logmeanexp(A; dims)`. +All three functions have +[ChainRules](https://github.com/JuliaDiff/ChainRulesCore.jl) derivative rules +(loaded automatically when ChainRulesCore is in the environment), so they can be +differentiated with ChainRules-based AD packages such as Zygote. + ## Reference ```@docs diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl new file mode 100644 index 0000000..98d1e03 --- /dev/null +++ b/ext/LogStatFunctionsChainRulesCoreExt.jl @@ -0,0 +1,74 @@ +module LogStatFunctionsChainRulesCoreExt + +using LogStatFunctions: logmeanexp, logvarexp, logstdexp +import ChainRulesCore + +function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :) + Ω = logmeanexp(x; dims) + n = length(x) ÷ length(Ω) + ΔΩ = sum(exp.(x .- Ω) .* Δx; dims) ./ n + return Ω, ΔΩ +end + +function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :) + Ω = logmeanexp(x; dims) + n = length(x) ÷ length(Ω) + return Ω, _∂x_pullback(exp.(x .- Ω) ./ n, x) +end + +function ChainRulesCore.frule( + (_, Δx), ::typeof(logvarexp), x::AbstractArray{<:Real}; + dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims) + ) + Ω = logvarexp(x; dims, corrected, logmean) + ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) .* Δx; dims) + return Ω, ΔΩ +end + +function ChainRulesCore.rrule( + ::typeof(logvarexp), x::AbstractArray{<:Real}; + dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims) + ) + Ω = logvarexp(x; dims, corrected, logmean) + return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims), x) +end + +function ChainRulesCore.frule( + (_, Δx), ::typeof(logstdexp), x::AbstractArray{<:Real}; + dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims) + ) + Ω = logstdexp(x; dims, corrected, logmean) + ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) ./ 2 .* Δx; dims) + return Ω, ΔΩ +end + +function ChainRulesCore.rrule( + ::typeof(logstdexp), x::AbstractArray{<:Real}; + dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims) + ) + Ω = logstdexp(x; dims, corrected, logmean) + return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims) / 2, x) +end + +# ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean). +# The m-dependence on x drops out because Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts +# the result by a constant, so the gradient is the same either way. +function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims) + d = x .- logmean + e = expm1.(d) + return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims) +end + +function _∂x_pullback(∂x, x) + project_x = ChainRulesCore.ProjectTo(x) + function pullback(Ω̄) + x̄ = ChainRulesCore.InplaceableThunk( + Δ -> Δ .+= Ω̄ .* ∂x, + ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), + ) + return ChainRulesCore.NoTangent(), x̄ + end + return pullback +end + +end # module diff --git a/test/Project.toml b/test/Project.toml index 9bc888c..f0d671d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" LogStatFunctions = "1b3c2205-da69-4893-b2df-717075b12251" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -8,6 +9,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Aqua = "0.8" +ChainRulesTestUtils = "1" ExplicitImports = "1.15" StaticArrays = "1" Statistics = "1.10" diff --git a/test/chainrules.jl b/test/chainrules.jl new file mode 100644 index 0000000..be46b76 --- /dev/null +++ b/test/chainrules.jl @@ -0,0 +1,17 @@ +using Test: @testset +using ChainRulesTestUtils: test_frule, test_rrule +using LogStatFunctions: logmeanexp, logvarexp, logstdexp + +@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 + test_frule(logmeanexp, x; fkwargs = (; dims)) + test_rrule(logmeanexp, x; fkwargs = (; dims)) + for corrected in (true, false) + test_frule(logvarexp, x; fkwargs = (; dims, corrected)) + test_rrule(logvarexp, x; fkwargs = (; dims, corrected)) + test_frule(logstdexp, x; fkwargs = (; dims, corrected)) + test_rrule(logstdexp, x; fkwargs = (; dims, corrected)) + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 84bfb2f..a8459ec 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,9 @@ module functions_tests include("functions.jl") end +module chainrules_tests + include("chainrules.jl") +end module explicit_imports_tests include("explicit_imports.jl") end