diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 51c3949ae9..0fb2ad6ec8 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -176,6 +176,7 @@ Extensions are in `ext/` directory: - EnzymeBFloat16sExt - EnzymeChainRulesCoreExt - EnzymeGPUArraysCoreExt +- EnzymeGammaExt - EnzymeLogExpFunctionsExt - EnzymeSpecialFunctionsExt - EnzymeStaticArraysExt diff --git a/Project.toml b/Project.toml index bb6a495111..5d6be65ce1 100644 --- a/Project.toml +++ b/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" diff --git a/ext/EnzymeGammaExt.jl b/ext/EnzymeGammaExt.jl new file mode 100644 index 0000000000..adbe746f27 --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 60361f447b..6dd76aec98 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/ext/gamma.jl b/test/ext/gamma.jl new file mode 100644 index 0000000000..838baf0981 --- /dev/null +++ b/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