Add EnzymeGammaExt with analytic derivative rule for Gamma.gamma

.github/copilot-instructions.md

@@ -176,6 +176,7 @@ Extensions are in `ext/` directory:
 - EnzymeBFloat16sExt
 - EnzymeChainRulesCoreExt
 - EnzymeGPUArraysCoreExt
+- EnzymeGammaExt
 - EnzymeLogExpFunctionsExt
 - EnzymeSpecialFunctionsExt
 - EnzymeStaticArraysExt

Project.toml

@@ -24,6 +24,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+Gamma = "a0844989-3bd2-4988-8bea-c9407ab0941b"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -32,6 +33,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 EnzymeBFloat16sExt = "BFloat16s"
 EnzymeChainRulesCoreExt = "ChainRulesCore"
 EnzymeGPUArraysCoreExt = "GPUArraysCore"
+EnzymeGammaExt = "Gamma"
 EnzymeLogExpFunctionsExt = "LogExpFunctions"
 EnzymeSpecialFunctionsExt = "SpecialFunctions"
 EnzymeStaticArraysExt = "StaticArrays"
@@ -45,6 +47,7 @@ EnzymeCore = "0.8.21"
 Enzyme_jll = "0.0.279"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"
+Gamma = "1"
 LLVM = "9.10"
 LogExpFunctions = "0.3, 1"
 ObjectFile = "0.4, 0.5"
@@ -60,6 +63,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+Gamma = "a0844989-3bd2-4988-8bea-c9407ab0941b"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

ext/EnzymeGammaExt.jl

@@ -0,0 +1,19 @@
+module EnzymeGammaExt
+
+using Gamma
+using Enzyme
+
+# `Gamma.gamma` (used e.g. by HypergeometricFunctions >= 0.3.29, and thus by the
+# incomplete-beta / truncated-distribution code paths in Distributions) is a
+# pure-Julia Cephes implementation whose internal `while` loops Enzyme's reverse
+# mode cannot rematerialize when the call is inlined into a larger differentiated
+# function (it fails with an internal "ficticious phi" recompute error). Provide
+# the analytic derivative Γ'(x) = Γ(x) ψ(x) so Enzyme differentiates the closed
+# form instead of descending into the loops. `Ω` is the primal result Γ(x).
+Enzyme.EnzymeRules.@easy_rule(
+    Gamma.gamma(x::AbstractFloat),
+    @setup(),
+    (Ω * Gamma.digamma(x),),
+)
+
+end

test/Project.toml

@@ -9,6 +9,8 @@ FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 Enzyme_jll = "7cc45869-7501-5eee-bdea-0790c847d4ef"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+Gamma = "a0844989-3bd2-4988-8bea-c9407ab0941b"
+HypergeometricFunctions = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"

test/ext/gamma.jl

@@ -0,0 +1,40 @@
+using Gamma
+using HypergeometricFunctions
+using FiniteDifferences
+
+include("../common.jl")
+
+Enzyme.Compiler.VERBOSE_ERRORS[] = true
+
+# `Gamma.gamma` is a pure-Julia Cephes implementation with internal `while` loops.
+# Differentiating it in isolation works, but when it is inlined into a larger
+# reverse-mode call (e.g. via HypergeometricFunctions' ₂F₁, which the incomplete
+# beta / truncated Beta path in Distributions relies on) Enzyme could not
+# rematerialize the loop and failed with an internal "ficticious phi" error.
+# EnzymeGammaExt supplies the analytic rule Γ'(x) = Γ(x) ψ(x); these tests check
+# both the standalone derivative and the previously-failing in-context case.
+
+@testset "Gamma.gamma scalar derivative" begin
+    # Points chosen to exercise every branch of the Cephes implementation:
+    # the x<2 upshift loop, the plain 2≤x<3 region, the x≥3 downshift loop,
+    # the x>11.5 asymptotic branch, and the x<0 reflection formula.
+    for x in (0.5, 1.5, 2.5, 4.5, 8.5, 12.5, -0.6, -2.6)
+        test_scalar(Gamma.gamma, x; rtol = 1.0e-5, atol = 1.0e-5)
+    end
+    test_scalar(Gamma.gamma, 0.5f0; rtol = 1.0e-4, atol = 1.0e-4)
+    test_scalar(Gamma.gamma, 4.5f0; rtol = 1.0e-4, atol = 1.0e-4)
+end
+
+@testset "Gamma.gamma in-context (HypergeometricFunctions ₂F₁)" begin
+    # Reverse mode differentiating ₂F₁ w.r.t. all of (a, b, c) routes through the
+    # parameter-derivative of `gamma` and reproduced the original failure.
+    F(a, b, c, z) = HypergeometricFunctions._₂F₁(a, b, c, z)
+    a, b, c, z = 1.1, 1.3, 2.3, 0.4
+    fd = collect(FiniteDifferences.grad(central_fdm(5, 1), F, a, b, c, z))
+
+    rev = collect(Float64.(Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(F), a, b, c, z)))
+    @test isapprox(rev, fd; rtol = 1.0e-5, atol = 1.0e-5)
+
+    fwd = collect(Float64.(Enzyme.gradient(Enzyme.Forward, Enzyme.Const(F), a, b, c, z)))
+    @test isapprox(fwd, fd; rtol = 1.0e-5, atol = 1.0e-5)
+end