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.288"
 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,12 @@
+module EnzymeGammaExt
+
+using Gamma
+using Enzyme
+
+Enzyme.EnzymeRules.@easy_rule(
+    Gamma.gamma(x::AbstractFloat),
+    @setup(),
+    (Ω * Gamma.digamma(x),),
+)
+
+end

test/Project.toml

@@ -10,6 +10,7 @@ FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+Gamma = "a0844989-3bd2-4988-8bea-c9407ab0941b"
 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,31 @@
+using Gamma
+using FiniteDifferences
+
+include("../common.jl")
+
+Enzyme.Compiler.VERBOSE_ERRORS[] = true
+
+@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 (beta function)" begin
+    # The rule must also fire for `gamma` calls inside a larger differentiated
+    # function, with partials accumulating across several call sites.
+    B(a, b) = Gamma.gamma(a) * Gamma.gamma(b) / Gamma.gamma(a + b)
+    a, b = 1.1, 2.3
+    fd = collect(FiniteDifferences.grad(central_fdm(5, 1), B, a, b))
+
+    rev = collect(Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(B), a, b))
+    @test isapprox(rev, fd; rtol = 1.0e-6, atol = 1.0e-6)
+
+    fwd = collect(Enzyme.gradient(Enzyme.Forward, Enzyme.Const(B), a, b))
+    @test isapprox(fwd, fd; rtol = 1.0e-6, atol = 1.0e-6)
+end