Add rrule for logpdf of NegativeBinomial (completes #1568) (#1579)

* Add `rrule` for `logpdf` of `NegativeBinomial`

* Remove unnecessary module prefix

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Use explicit division

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Refator and correct pullback

* Add tests for rrule

* Use `forward_fdm` for testing `rrule`

* Fix tests

* Update test/negativebinomial.jl

* Fix tests (without `p = 1 - eps()`)

* Use FD for all tests, use random parameters

* Avoid type instability

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Split and rearrange ForwardDiff and rrule tests

* Bump version

* Fix typo

* Clean tests (revert unrelated changes) and fix them

Co-authored-by: Simone Carlo Surace <51025924+simsurace@users.noreply.github.com>
Co-authored-by: Simone Carlo Surace <simone.surace@artificialy.com>

Author: 3 people

src/univariate/discrete/negativebinomial.jl

@@ -135,3 +135,38 @@ function cf(d::NegativeBinomial, t::Real)
     r, p = params(d)
     return (((1 - p) * cis(t)) / (1 - p * cis(t)))^r
 end
+
+# ChainRules definitions
+
+function ChainRulesCore.rrule(::typeof(logpdf), d::NegativeBinomial, k::Real)
+    # Compute log probability
+    r, p = params(d)
+    edgecase = isone(p) && iszero(k)
+    insupp = insupport(d, k)
+    
+    # Primal computation
+    Ω = r * log(p) + k * log1p(-p)
+    if edgecase
+        Ω = zero(Ω)
+    elseif !insupp
+        Ω = oftype(Ω, -Inf)
+    else
+        Ω = Ω - log(k + r) - logbeta(r, k + 1)
+    end
+
+    # Define pullback
+    function logpdf_NegativeBinomial_pullback(Δ)
+        Δr = Δ * (log(p) - inv(k + r) - digamma(r) + digamma(r + k + 1))
+        Δp = Δ * (r / p - k / (1 - p))
+        if edgecase
+            Δp = oftype(Δp, Δ * r)
+        elseif !insupp
+            Δr = oftype(Δr, NaN)
+            Δp = oftype(Δp, NaN)
+        end
+        Δd = ChainRulesCore.Tangent{typeof(d)}(; r=Δr, p=Δp)
+        return ChainRulesCore.NoTangent(), Δd, ChainRulesCore.NoTangent()
+    end
+
+    return Ω, logpdf_NegativeBinomial_pullback
+end

test/negativebinomial.jl

@@ -1,7 +1,9 @@
 using Distributions
 using Test, ForwardDiff
+using ChainRulesTestUtils
+using FiniteDifferences
 
-# Currently, most of the tests for NegativeBinomail are in the "ref" folder.
+# Currently, most of the tests for NegativeBinomial are in the "ref" folder.
 # Eventually, we might want to consolidate the tests here
 
 mydiffp(r, p, k) = r/p - k/(1 - p)
@@ -19,3 +21,25 @@ end
     @test logpdf(NegativeBinomial(0.5, 1.0), 1) === -Inf
     @test all(iszero, rand(NegativeBinomial(rand(), 1.0), 10))
 end
+
+@testset "rrule: logpdf of NegativeBinomial" begin
+    r = randexp()
+
+    # Test with values in and outside of support
+    p = rand()
+    dist = NegativeBinomial(r, p)
+    fdm = central_fdm(5, 1; max_range=min(r, p, 1-p)/2) # avoids numerical issues with finite differencing
+    for k in (0, 10, 42, -1, -5, -13)
+        # Test both integers and floating point numbers.
+        # For floating point numbers we have to tell FiniteDifferences explicitly that the
+        # argument is non-differentiable. Otherwise it will compute `NaN` as derivative.
+        test_rrule(logpdf, dist, k; fdm=fdm, nans=true)
+        test_rrule(logpdf, dist, float(k) ⊢ ChainRulesTestUtils.NoTangent(); fdm=fdm, nans=true)
+    end
+
+    # Test edge case `p = 1` and `k = 0`
+    dist = NegativeBinomial(r, 1)
+    fdm = backward_fdm(5, 1; max_range = r/10)
+    test_rrule(logpdf, dist, 0; fdm=fdm)
+    test_rrule(logpdf, dist, 0.0 ⊢ ChainRulesTestUtils.NoTangent(); fdm=fdm)
+end