Compute logvarexp/logstdexp gradients in the log domain

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.

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
 

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

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"

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