From ad5d8c7794de6ab6fea2b04046773b5e150bded9 Mon Sep 17 00:00:00 2001 From: mohamed82008 Date: Sun, 27 Jan 2019 15:27:33 +1100 Subject: [PATCH 1/4] ad workaround for nbinomlogpdf --- src/core/ad.jl | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/core/ad.jl b/src/core/ad.jl index f3b07cbf1..d32ab8bb8 100644 --- a/src/core/ad.jl +++ b/src/core/ad.jl @@ -180,6 +180,12 @@ Tracker.@grad function binomlogpdf(n::Int, p::Tracker.TrackedReal, x::Int) Δ->(nothing, Δ * (x / p - (n - x) / (1 - p)), nothing) end +import StatsFuns: nbinomlogpdf +nbinomlogpdf(n::Tracker.TrackedReal, p::Tracker.TrackedReal, x::Int) = Tracker.track(nbinomlogpdf, n, p, x) +Tracker.@grad function nbinomlogpdf(r::Tracker.TrackedReal, p::Tracker.TrackedReal, k::Int) + return nbinomlogpdf(Tracker.data(r), Tracker.data(p), k), + Δ->(Δ * (sum(1 / (k + r - i) for i in 1:k) + log(1 - p)), Δ * (-r / (1 - p) + k / p), nothing) +end import StatsFuns: poislogpdf poislogpdf(v::Tracker.TrackedReal, x::Int) = Tracker.track(poislogpdf, v, x) @@ -195,6 +201,17 @@ function binomlogpdf(n::Int, p::ForwardDiff.Dual{T}, x::Int) where {T} return FD(binomlogpdf(n, val, x), Δ * (x / val - (n - x) / (1 - val))) end +function nbinomlogpdf(r::ForwardDiff.Dual{T}, p::ForwardDiff.Dual{T}, k::Int) where {T} + FD = ForwardDiff.Dual{T} + val_p = ForwardDiff.value(p) + val_r = ForwardDiff.value(r) + + Δ_p = ForwardDiff.partials(p) * (-val_r / (1 - val_p) + k / val_p) + Δ_r = ForwardDiff.partials(r) * (sum(1 / (k + val_r - i) for i in 1:k) + log(1 - val_p)) + Δ = Δ_p + Δ_r + return FD(nbinomlogpdf(val_r, val_p, k), Δ) +end + function poislogpdf(v::ForwardDiff.Dual{T}, x::Int) where {T} FD = ForwardDiff.Dual{T} val = ForwardDiff.value(v) From 3aca27ef942aa9a12bdfff95a18595a3a635d838 Mon Sep 17 00:00:00 2001 From: mohamed82008 Date: Sun, 3 Feb 2019 11:35:27 +1100 Subject: [PATCH 2/4] fix nbinomlogpdf derivatives --- src/core/ad.jl | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/core/ad.jl b/src/core/ad.jl index d32ab8bb8..84e6e9881 100644 --- a/src/core/ad.jl +++ b/src/core/ad.jl @@ -181,10 +181,26 @@ Tracker.@grad function binomlogpdf(n::Int, p::Tracker.TrackedReal, x::Int) end import StatsFuns: nbinomlogpdf +# Note the definition of NegativeBinomial in Julia is not the same as Wikipedia's. +# Check the docstring of NegativeBinomial, r is the number of successes and +# k is the number of failures +_nbinomlogpdf_grad_1(r, p, k) = sum(1 / (k + r - i) for i in 1:k) + log(p) +_nbinomlogpdf_grad_2(r, p, k) = -k / (1 - p) + r / p + nbinomlogpdf(n::Tracker.TrackedReal, p::Tracker.TrackedReal, x::Int) = Tracker.track(nbinomlogpdf, n, p, x) +nbinomlogpdf(n::Real, p::Tracker.TrackedReal, x::Int) = Tracker.track(nbinomlogpdf, n, p, x) +nbinomlogpdf(n::Tracker.TrackedReal, p::Real, x::Int) = Tracker.track(nbinomlogpdf, n, p, x) Tracker.@grad function nbinomlogpdf(r::Tracker.TrackedReal, p::Tracker.TrackedReal, k::Int) return nbinomlogpdf(Tracker.data(r), Tracker.data(p), k), - Δ->(Δ * (sum(1 / (k + r - i) for i in 1:k) + log(1 - p)), Δ * (-r / (1 - p) + k / p), nothing) + Δ->(Δ * _nbinomlogpdf_grad_1(r, p, k), Δ * _nbinomlogpdf_grad_2(r, p, k), nothing) +end +Tracker.@grad function nbinomlogpdf(r::Real, p::Tracker.TrackedReal, k::Int) + return nbinomlogpdf(Tracker.data(r), Tracker.data(p), k), + Δ->(Tracker._zero(r), Δ * _nbinomlogpdf_grad_2(r, p, k), nothing) +end +Tracker.@grad function nbinomlogpdf(r::Tracker.TrackedReal, p::Real, k::Int) + return nbinomlogpdf(Tracker.data(r), Tracker.data(p), k), + Δ->(Δ * _nbinomlogpdf_grad_1(r, p, k), Tracker._zero(p), nothing) end import StatsFuns: poislogpdf @@ -206,11 +222,23 @@ function nbinomlogpdf(r::ForwardDiff.Dual{T}, p::ForwardDiff.Dual{T}, k::Int) wh val_p = ForwardDiff.value(p) val_r = ForwardDiff.value(r) - Δ_p = ForwardDiff.partials(p) * (-val_r / (1 - val_p) + k / val_p) - Δ_r = ForwardDiff.partials(r) * (sum(1 / (k + val_r - i) for i in 1:k) + log(1 - val_p)) + Δ_r = ForwardDiff.partials(r) * _nbinomlogpdf_grad_1(val_r, val_p, k) + Δ_p = ForwardDiff.partials(p) * _nbinomlogpdf_grad_2(val_r, val_p, k) Δ = Δ_p + Δ_r return FD(nbinomlogpdf(val_r, val_p, k), Δ) end +function nbinomlogpdf(r::Real, p::ForwardDiff.Dual{T}, k::Int) where {T} + FD = ForwardDiff.Dual{T} + val_p = ForwardDiff.value(p) + Δ_p = ForwardDiff.partials(p) * _nbinomlogpdf_grad_2(r, val_p, k) + return FD(nbinomlogpdf(r, val_p, k), Δ_p) +end +function nbinomlogpdf(r::ForwardDiff.Dual{T}, p::Real, k::Int) where {T} + FD = ForwardDiff.Dual{T} + val_r = ForwardDiff.value(r) + Δ_r = ForwardDiff.partials(r) * _nbinomlogpdf_grad_1(val_r, p, k) + return FD(nbinomlogpdf(val_r, p, k), Δ_r) +end function poislogpdf(v::ForwardDiff.Dual{T}, x::Int) where {T} FD = ForwardDiff.Dual{T} From 5d4898f38a04fba004382d6809ff67d0df6b3b02 Mon Sep 17 00:00:00 2001 From: mohamed82008 Date: Sun, 3 Feb 2019 11:35:46 +1100 Subject: [PATCH 3/4] add negative binomial logpdf AD tests --- .../AD_compatibility_with_distributions.jl | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/test/ad.jl/AD_compatibility_with_distributions.jl b/test/ad.jl/AD_compatibility_with_distributions.jl index 9b3e3b7fe..4215d3c08 100644 --- a/test/ad.jl/AD_compatibility_with_distributions.jl +++ b/test/ad.jl/AD_compatibility_with_distributions.jl @@ -88,3 +88,63 @@ let atol=1e-8, ) end + +let + foo = p->Turing.nbinomlogpdf(5, p, 1) + @test isapprox( + Tracker.gradient(foo, 0.5)[1], + central_fdm(5, 1)(foo, 0.5); + rtol=1e-8, + atol=1e-8, + ) + @test isapprox( + Tracker.gradient(foo, 0.5)[1], + ForwardDiff.derivative(foo, 0.5); + rtol=1e-8, + atol=1e-8, + ) + + bar = p->logpdf(NegativeBinomial(5, p), 3) + @test isapprox( + Tracker.gradient(bar, 0.5)[1], + central_fdm(5, 1)(bar, 0.5); + rtol=1e-8, + atol=1e-8, + ) + @test isapprox( + Tracker.gradient(bar, 0.5)[1], + ForwardDiff.derivative(bar, 0.5); + rtol=1e-8, + atol=1e-8, + ) +end + +let + foo = r->Turing.nbinomlogpdf(r, 0.5, 1) + @test isapprox( + Tracker.gradient(foo, 3.5)[1], + central_fdm(5, 1)(foo, 3.5); + rtol=1e-8, + atol=1e-8, + ) + @test isapprox( + Tracker.gradient(foo, 3.5)[1], + ForwardDiff.derivative(foo, 3.5); + rtol=1e-8, + atol=1e-8, + ) + + bar = r->logpdf(NegativeBinomial(r, 0.5), 3) + @test isapprox( + Tracker.gradient(bar, 3.5)[1], + central_fdm(5, 1)(bar, 3.5); + rtol=1e-8, + atol=1e-8, + ) + @test isapprox( + Tracker.gradient(bar, 3.5)[1], + ForwardDiff.derivative(bar, 3.5); + rtol=1e-8, + atol=1e-8, + ) +end From aea53cdc81d8e230828439539747f005506aca7b Mon Sep 17 00:00:00 2001 From: mohamed82008 Date: Sun, 3 Feb 2019 11:51:26 +1100 Subject: [PATCH 4/4] add a test for the full gradient of nbinomlogpdf --- test/ad.jl/AD_compatibility_with_distributions.jl | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/test/ad.jl/AD_compatibility_with_distributions.jl b/test/ad.jl/AD_compatibility_with_distributions.jl index 4215d3c08..b1198bd8e 100644 --- a/test/ad.jl/AD_compatibility_with_distributions.jl +++ b/test/ad.jl/AD_compatibility_with_distributions.jl @@ -148,3 +148,13 @@ let atol=1e-8, ) end + +let + foo = x -> Turing.nbinomlogpdf(x[1], x[2], 1) + @test isapprox( + Tracker.gradient(foo, [3.5, 0.5])[1], + ForwardDiff.gradient(foo, [3.5, 0.5]); + rtol=1e-8, + atol=1e-8, + ) +end