AD workaround for nbinomlogpdf (#664)

* ad workaround for nbinomlogpdf

* fix nbinomlogpdf derivatives

* add negative binomial logpdf AD tests

* add a test for the full gradient of nbinomlogpdf

src/core/ad.jl

@@ -180,6 +180,28 @@ Tracker.@grad function binomlogpdf(n::Int, p::Tracker.TrackedReal, x::Int)
  Δ->(nothing, Δ * (x / p - (n - x) / (1 - p)), nothing)
 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),
+ Δ->(Δ * _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
 poislogpdf(v::Tracker.TrackedReal, x::Int) = Tracker.track(poislogpdf, v, x)
@@ -195,6 +217,29 @@ 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)
+
+ Δ_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}
  val = ForwardDiff.value(v)

test/ad.jl/AD_compatibility_with_distributions.jl

@@ -88,3 +88,73 @@ 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
+
+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