logpdf test

Author: matbesancon

src/multivariate/dirichlet.jl

@@ -381,7 +381,6 @@ function fit_mle(::Type{<:Dirichlet}, P::AbstractMatrix{Float64},
 end
 
 ## Differentiation
-using Test
 function ChainRulesCore.frule((_, Δalpha), DT::Union{Type{Dirichlet{T}}, Type{Dirichlet}}, alpha::AbstractVector{T}; check_args = true) where {T}
     d = DT(alpha; check_args=check_args)
     Δalpha = ChainRulesCore.unthunk(Δalpha)
@@ -401,3 +400,42 @@ function ChainRulesCore.rrule(DT::Union{Type{Dirichlet{T}}, Type{Dirichlet}}, al
     end
     return d, dirichlet_pullback
 end
+
+function ChainRulesCore.frule((_, Δd, Δx), ::typeof(_logpdf), d::Dirichlet, x::AbstractVector{T}) where {T}
+    lp = _logpdf(d, x)
+    if !insupport(d, x)
+        return (lp, zero(lp))
+    end
+    ∂α = sum(Δd.alpha[i] * log(x[i]) for i in eachindex(x))
+    ∂l = - Δd.lmnB
+    ∂x = sum((d.alpha[i] - 1) * Δx[i] / x[i] for i in eachindex(x))
+    return (lp, ∂α + ∂l + ∂x)
+end
+
+function ChainRulesCore.rrule(::typeof(_logpdf), d::Dirichlet, x::AbstractVector{T}) where {T}
+    y = _logpdf(d, x)
+    function Dirichlet_logpdf_pullback(dy)
+        if !isfinite(y)
+            backing = (alpha = zero(d.alpha), alpha0 = ChainRulesCore.ZeroTangent(), lmnB=zero(d.lmnB))
+            ∂d = ChainRulesCore.Tangent{typeof(d), typeof(backing)}(backing)
+            ∂x = zero(d.alpha + x)
+            return (ChainRulesCore.NoTangent(), ∂d, ∂x)
+        end
+        ∂alpha = dy * log.(x)
+        ∂l = -dy
+        ∂x = dy * (d.alpha .-1) ./ x
+        backing = (alpha = ∂alpha, alpha0 = ChainRulesCore.ZeroTangent(), lmnB=∂l)
+        ∂d = ChainRulesCore.Tangent{typeof(d), typeof(backing)}(backing)
+        return (ChainRulesCore.NoTangent(), ∂d, ∂x)
+    end
+    return (y, Dirichlet_logpdf_pullback)
+end
+
+function _logpdf(d::Dirichlet, x::AbstractVector{<:Real})
+    if !insupport(d, x)
+        return xlogy(one(eltype(d.alpha)), zero(eltype(x))) - d.lmnB
+    end
+    a = d.alpha
+    s = sum(xlogy(αi - 1, xi) for (αi, xi) in zip(d.alpha, x))
+    return s - d.lmnB
+end

test/dirichlet.jl

@@ -129,14 +129,53 @@ end
     @test entropy(Dirichlet(ones(N))) ≈ -loggamma(N)
 end
 
-@testset "Dirichlet differentiation" begin
-    for n in (2, 10)
-        alpha = rand(n)
-        Δalpha = randn(n)
-        d2, ∂d = ChainRulesCore.frule((nothing, Δalpha), Dirichlet, alpha)
-        ChainRulesTestUtils.test_frule(Dirichlet ⊢ ChainRulesCore.NoTangent(), alpha ⊢ Δalpha, check_inferred=true)
-
-        _, dp = ChainRulesCore.rrule(Dirichlet, alpha)
-        ChainRulesTestUtils.test_rrule(Dirichlet{Float64} ⊢ ChainRulesCore.NoTangent(), alpha)
+@testset "Dirichlet differentiation $n" for n in (2, 10)
+    alpha = rand(n)
+    Δalpha = randn(n)
+    d, ∂d = ChainRulesCore.frule((nothing, Δalpha), Dirichlet, alpha)
+    ChainRulesTestUtils.test_frule(Dirichlet ⊢ ChainRulesCore.NoTangent(), alpha ⊢ Δalpha)
+    _, dp = ChainRulesCore.rrule(Dirichlet, alpha)
+    ChainRulesTestUtils.test_rrule(Dirichlet{Float64} ⊢ ChainRulesCore.NoTangent(), alpha)
+    x = rand(n)
+    x ./= sum(x)
+    Δx = 0.05 * rand(n)
+    Δx .-= mean(Δx)
+    # such that x ∈ Δ, x + Δx ∈ Δ
+    ChainRulesTestUtils.test_frule(Distributions._logpdf ⊢ ChainRulesCore.NoTangent(), d, x ⊢ Δx)
+    @testset "finite diff f/r-rule logpdf" begin
+        for _ in 1:10
+            x = rand(n)
+            x ./= sum(x)
+            Δx = 0.005 * rand(n)
+            Δx .-= mean(Δx)
+            if insupport(d, x + Δx) && insupport(d, x - Δx)
+                y, pullback = ChainRulesCore.rrule(Distributions._logpdf, d, x)
+                yf, Δy = ChainRulesCore.frule(
+                    (
+                        ChainRulesCore.NoTangent(),
+                        map(zero, ChainRulesTestUtils.rand_tangent(d)),
+                        Δx,
+                    ),
+                    Distributions._logpdf,
+                    d, x,
+                )
+                y2 = Distributions._logpdf(d, x + Δx)
+                y1 = Distributions._logpdf(d, x - Δx)
+                @test isfinite(y)
+                @test y == yf
+                @test Δy ≈ y2 - y atol=5e-3
+                _, ∂d, ∂x = pullback(1.0)
+                @test y2 - y1 ≈ dot(2Δx, ∂x) atol=5e-3 rtol=1e-6
+                # mutating alpha only to compute a new y, changing only this term and not the others in Dirichlet
+                Δalpha = 0.03 * rand(n)
+                Δalpha .-= mean(Δalpha)
+                @assert all(>=(0), alpha + Δalpha)
+                d.alpha .+= Δalpha
+                ya = Distributions._logpdf(d, x)
+                # resetting alpha
+                d.alpha .-= Δalpha
+                @test ya - y ≈ dot(Δalpha, ∂d.alpha) atol=5e-5 rtol=1e-6
+            end
+        end
     end
 end