Merge pull request #305 from devmotion/dw/chainrules_gamma

Add ChainRules definitions for `gamma(a, x)`, `loggamma(a, x)`, and `gamma_inc`

Author: andreasnoack

src/chainrules.jl

@@ -3,6 +3,12 @@ derivatives of Bessel functions with respect to the order are not implemented cu
 https://github.com/JuliaMath/SpecialFunctions.jl/issues/160
 """
 
+const INCOMPLETE_GAMMA_INFO = """
+derivatives of the incomplete Gamma functions with respect to parameter `a` are not
+implemented currently:
+https://github.com/JuliaMath/SpecialFunctions.jl/issues/317
+"""
+
 ChainRulesCore.@scalar_rule(airyai(x), airyaiprime(x))
 ChainRulesCore.@scalar_rule(airyaiprime(x), x * airyai(x))
 ChainRulesCore.@scalar_rule(airybi(x), airybiprime(x))
@@ -26,6 +32,27 @@ ChainRulesCore.@scalar_rule(erfcx(x), (2 * x * Ω) - (2 / sqrt(π)))
 ChainRulesCore.@scalar_rule(erfi(x), (2 / sqrt(π)) * exp(x * x))
 ChainRulesCore.@scalar_rule(erfinv(x), (sqrt(π) / 2) * exp(Ω^2))
 ChainRulesCore.@scalar_rule(gamma(x), Ω * digamma(x))
+ChainRulesCore.@scalar_rule(
+    gamma(a, x),
+    (
+        ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO),
+        - exp(-x) * x^(a - 1),
+    ),
+)
+ChainRulesCore.@scalar_rule(
+    gamma_inc(a, x, IND),
+    @setup(z = exp(-x) * x^(a - 1) / gamma(a)),
+    (
+        ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO),
+        z,
+        ChainRulesCore.DoesNotExist(),
+    ),
+    (
+        ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO),
+        -z,
+        ChainRulesCore.DoesNotExist(),
+    ),
+)
 ChainRulesCore.@scalar_rule(
     invdigamma(x),
     inv(trigamma(invdigamma(x))),
@@ -98,3 +125,10 @@ ChainRulesCore.@scalar_rule(
 ChainRulesCore.@scalar_rule(logabsgamma(x), digamma(x), ChainRulesCore.Zero())
 
 ChainRulesCore.@scalar_rule(loggamma(x), digamma(x))
+ChainRulesCore.@scalar_rule(
+    loggamma(a, x),
+    (
+        ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO),
+        -exp(- (x + Ω)) * x^(a - 1),
+    )
+)

test/chainrules.jl

@@ -81,18 +81,33 @@
 
     @testset "log gamma and co" begin
         # It is important that we have negative numbers with both odd and even integer parts
-        for x in (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
+        test_points = (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
+        for x in test_points
             for m in (0, 1, 2, 3)
                 test_frule(polygamma, m, x)
                 test_rrule(polygamma, m, x)
             end
 
             isreal(x) && x < 0 && continue
             test_scalar(loggamma, x)
+            for a in test_points
+                # ensure all complex if any complex for FiniteDifferences
+                _a, _x = promote(a, x)
+                test_frule(gamma, _a, _x; rtol=1e-8)
+                test_rrule(gamma, _a, _x; rtol=1e-8)
+
+                test_frule(loggamma, _a, _x)
+                test_rrule(loggamma, _a, _x)
+            end
 
             isreal(x) || continue
             test_frule(logabsgamma, x)
             test_rrule(logabsgamma, x; output_tangent=(randn(), randn()))
+            for a in test_points
+                isreal(a) && a > 0 || continue
+                test_frule(gamma_inc, a, x, 0)
+                test_rrule(gamma_inc, a, x, 0; output_tangent=(randn(), randn()))
+            end
         end
     end
 end