Add ChainRulesCore extension with frule/rrule for logmeanexp, logvarexp, logstdexp

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.

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"

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

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

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"

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

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