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 c96aeb1824..a0cb0d7f62 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.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" diff --git a/ext/EnzymeGammaExt.jl b/ext/EnzymeGammaExt.jl new file mode 100644 index 0000000000..b6b7802d17 --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 0203d4f6d1..6beeb49b38 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/ext/gamma.jl b/test/ext/gamma.jl new file mode 100644 index 0000000000..93102e24ae --- /dev/null +++ b/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