diff --git a/CHANGELOG.md b/CHANGELOG.md index 16650fe..2b07e06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,8 +5,8 @@ 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). + `logmeanexp`, `logvarexp` and `logstdexp` on real arrays, 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/docs/src/index.md b/docs/src/index.md index 874a37d..3a0acf0 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -30,8 +30,9 @@ logstdexp(A; corrected=false) 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. +for real arrays (loaded automatically when ChainRulesCore is in the +environment), so they can be differentiated with ChainRules-based AD packages +such as Zygote. Complex arrays are not covered by these rules. ## Reference diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl index 98d1e03..e013f9d 100644 --- a/ext/LogStatFunctionsChainRulesCoreExt.jl +++ b/ext/LogStatFunctionsChainRulesCoreExt.jl @@ -1,19 +1,18 @@ module LogStatFunctionsChainRulesCoreExt using LogStatFunctions: logmeanexp, logvarexp, logstdexp +using LogExpFunctions: logsumexp, softmax 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 + ΔΩ = sum(softmax(x; dims) .* Δx; dims) 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) + return Ω, _∂x_pullback(softmax(x; dims), x) end function ChainRulesCore.frule( @@ -50,21 +49,23 @@ function ChainRulesCore.rrule( 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. +# ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean); +# the dependence of m on x drops out since Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts +# the result by a constant. Log-domain to avoid under/overflow of the squared terms. function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims) d = x .- logmean - e = expm1.(d) - return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims) + l = log.(abs.(expm1.(d))) + S = logsumexp(2 .* l; dims) + return sign.(d) .* 2 .* exp.(d .+ l .- S) end function _∂x_pullback(∂x, x) project_x = ChainRulesCore.ProjectTo(x) function pullback(Ω̄) + ΔΩ = ChainRulesCore.unthunk(Ω̄) x̄ = ChainRulesCore.InplaceableThunk( - Δ -> Δ .+= Ω̄ .* ∂x, - ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)), + Δ -> Δ .+= ΔΩ .* ∂x, + ChainRulesCore.@thunk(project_x(ΔΩ .* ∂x)), ) return ChainRulesCore.NoTangent(), x̄ end diff --git a/test/Project.toml b/test/Project.toml index f0d671d..3d5d33f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" LogStatFunctions = "1b3c2205-da69-4893-b2df-717075b12251" @@ -9,6 +10,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Aqua = "0.8" +ChainRulesCore = "1" ChainRulesTestUtils = "1" ExplicitImports = "1.15" StaticArrays = "1" diff --git a/test/chainrules.jl b/test/chainrules.jl index be46b76..ee19ea6 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -1,5 +1,6 @@ -using Test: @testset +using Test: @test, @testset using ChainRulesTestUtils: test_frule, test_rrule +using ChainRulesCore: NoTangent, frule, rrule, unthunk using LogStatFunctions: logmeanexp, logvarexp, logstdexp @testset "chainrules" begin @@ -15,3 +16,28 @@ using LogStatFunctions: logmeanexp, logvarexp, logstdexp end end end + +@testset "chainrules extreme values" begin + # Nearly equal entries: the primal and its gradient are representable, but squaring + # exp(xᵢ) - m outside the log domain underflows and used to yield Inf/NaN gradients. + x = [-1.0e-200, 1.0e-200] + for (f, h) in ((logvarexp, 1), (logstdexp, 2)) + Ω, pb = rrule(f, x) + @test isfinite(Ω) + x̄ = unthunk(pb(1.0)[2]) + @test collect(x̄) ≈ [-1.0e200, 1.0e200] ./ h + Ω, ΔΩ = frule((NoTangent(), [1.0, 0.0]), f, x) + @test isfinite(Ω) + @test ΔΩ ≈ -1.0e200 / h + @test frule((NoTangent(), [0.0, 1.0]), f, x)[2] ≈ 1.0e200 / h + end +end + +@testset "chainrules abstract eltype" begin + x = Real[0.0, 1.0] + for f in (logmeanexp, logvarexp, logstdexp) + @test unthunk(rrule(f, x)[2](1.0)[2]) ≈ unthunk(rrule(f, [0.0, 1.0])[2](1.0)[2]) + @test frule((NoTangent(), [1.0, 0.0]), f, x)[2] ≈ + frule((NoTangent(), [1.0, 0.0]), f, [0.0, 1.0])[2] + end +end