conflict

Author: matbesancon

src/multivariate/dirichlet.jl

@@ -377,17 +377,17 @@ function fit_mle(::Type{<:Dirichlet}, P::AbstractMatrix{Float64},
 end
 
 ## Differentiation
-function ChainRulesCore.frule((_, Δalpha)::Tuple{Any,Any}, DT::Union{Type{Dirichlet{T}}, Type{Dirichlet}}, alpha::AbstractVector{T}; check_args = true) where {T}
+function ChainRulesCore.frule((_, Δalpha), ::Type{DT}, alpha::AbstractVector{T}; check_args::Bool = true) where {T <: Real, DT <: Union{Dirichlet{T}, Dirichlet}}
     d = DT(alpha; check_args=check_args)
     Δalpha = ChainRulesCore.unthunk(Δalpha)
     ∂alpha0 = sum(Δalpha)
-    ∂lmnB::typeof(∂alpha0) = sum(Δalpha[i] * (SpecialFunctions.digamma(alpha[i]) - SpecialFunctions.digamma(d.alpha0)) for i in eachindex(alpha))
+    ∂lmnB = sum(Δalpha[i] * (SpecialFunctions.digamma(alpha[i]) - SpecialFunctions.digamma(d.alpha0)) for i in eachindex(alpha))
     backing = (alpha=Δalpha, alpha0=∂alpha0, lmnB=∂lmnB)
     t = ChainRulesCore.Tangent{typeof(d), NamedTuple{(:alpha, :alpha0, :lmnB), Tuple{typeof(alpha), typeof(d.alpha0), typeof(d.lmnB)}}}(backing)
     return d, t
 end
 
-function ChainRulesCore.rrule(DT::Union{Type{Dirichlet{T}}, Type{Dirichlet}}, alpha::AbstractVector{T}; check_args = true) where {T}
+function ChainRulesCore.rrule(::Type{DT}, alpha::AbstractVector{T}; check_args::Bool = true) where {T <: Real, DT <: Union{Dirichlet{T}, Dirichlet}}
     d = DT(alpha; check_args=check_args)
     function dirichlet_pullback(d_dir)
         d_dir = ChainRulesCore.unthunk(d_dir)
@@ -426,12 +426,3 @@ function ChainRulesCore.rrule(::typeof(_logpdf), d::Dirichlet, x::AbstractVector
     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

@@ -4,6 +4,7 @@ using  Distributions
 using Test, Random, LinearAlgebra
 using ChainRulesCore
 using ChainRulesTestUtils
+using FiniteDifferences
 
 Random.seed!(34567)
 
@@ -132,7 +133,7 @@ end
 @testset "Dirichlet differentiation $n" for n in (2, 10)
     alpha = rand(n)
     Δalpha = randn(n)
-    d, ∂d = ChainRulesCore.frule((nothing, Δalpha), Dirichlet, alpha)
+    d, ∂d = @inferred ChainRulesCore.frule((nothing, Δalpha), Dirichlet, alpha)
     ChainRulesTestUtils.test_frule(Dirichlet ⊢ ChainRulesCore.NoTangent(), alpha ⊢ Δalpha; fdm=FiniteDifferences.forward_fdm(5, 1))
     _, dp = ChainRulesCore.rrule(Dirichlet, alpha)
     ChainRulesTestUtils.test_rrule(Dirichlet{Float64}, alpha)