Define diff rules as methods and move provider rules to extensions

.github/workflows/ci.yml

@@ -42,3 +42,19 @@ jobs:
         continue-on-error: true
         with:
           files: lcov.info
+
+  downgrade:
+    name: Downgrade
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: julia-actions/setup-julia@v3
+        with:
+          version: '1'
+      - uses: julia-actions/cache@v3
+      - uses: julia-actions/julia-downgrade-compat@v2
+        with:
+          projects: ., test
+          skip: Random,Test
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1

Project.toml

@@ -4,16 +4,22 @@ version = "1.17.0-dev"
 
 [deps]
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
+
+[weakdeps]
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
+[extensions]
+LogExpFunctionsExt = "LogExpFunctions"
+NaNMathExt = "NaNMath"
+NaNMathSpecialFunctionsExt = ["NaNMath", "SpecialFunctions"]
+SpecialFunctionsExt = "SpecialFunctions"
+
 [compat]
 IrrationalConstants = "0.1.1, 0.2"
 LogExpFunctions = "0.3.2, 1"
 NaNMath = "0.3, 1"
-Random = "1"
 SpecialFunctions = "1.1, 2"
 julia = "1.10"
 

docs/src/index.md

@@ -13,6 +13,20 @@ Note that DiffRules is *not* a fully-fledged symbolic differentiation tool. It i
 simple global database of common derivative definitions, and was developed with the goal of
 improving derivative coverage in downstream tools.
 
+Rules for SpecialFunctions, NaNMath and LogExpFunctions live in package extensions, so
+DiffRules itself depends on none of them. Load the package you need alongside DiffRules to
+get its rules:
+
+```julia
+using DiffRules, SpecialFunctions
+
+DiffRules.hasdiffrule(SpecialFunctions.erf, 1)  # true
+```
+
+Without `using SpecialFunctions`, that query returns `false` and `diffrules()` does not list
+its rules. Packages generating code from `diffrules()` should therefore load the packages
+whose rules they want, and skip rules for modules they do not have in scope.
+
 ```@docs
 DiffRules.@define_diffrule
 DiffRules.diffrule

ext/LogExpFunctionsExt.jl

@@ -0,0 +1,36 @@
+module LogExpFunctionsExt
+
+using DiffRules: @define_diffrule
+using LogExpFunctions:
+    LogExpFunctions, log1mexp, log2mexp, logaddexp, logexpm1, logistic, logsubexp
+
+###################
+# LogExpFunctions #
+###################
+
+# unary
+@define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x))
+@define_diffrule LogExpFunctions.logistic(x) = :(z = $logistic($x); z * (1 - z))
+@define_diffrule LogExpFunctions.logit(x) = :(inv($x * (1 - $x)))
+@define_diffrule LogExpFunctions.log1psq(x) = :(2 * $x / (1 + $x^2))
+@define_diffrule LogExpFunctions.log1pexp(x) = :($logistic($x))
+@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - $log1mexp($x)))
+@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - $log2mexp($x)))
+@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - $logexpm1($x)))
+@define_diffrule LogExpFunctions.log1pmx(x) = :(-$x / (1 + $x))
+@define_diffrule LogExpFunctions.logmxp1(x) = :((1 - $x) / $x)
+
+# binary
+@define_diffrule LogExpFunctions.xlogy(x, y) =
+    :(log($y)),
+    :(z = $x / $y; iszero($x) && !isnan($y) ? zero(z) : z)
+@define_diffrule LogExpFunctions.logaddexp(x, y) =
+    :(exp($x - $logaddexp($x, $y))), :(exp($y - $logaddexp($x, $y)))
+@define_diffrule LogExpFunctions.logsubexp(x, y) =
+    :(z = $logsubexp($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
+    :(z = $logsubexp($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
+@define_diffrule LogExpFunctions.xlog1py(x, y) =
+    :(log1p($y)),
+    :(z = $x / (1 + $y); iszero($x) && !isnan($y) ? zero(z) : z)
+
+end # module

ext/NaNMathExt.jl

@@ -0,0 +1,36 @@
+module NaNMathExt
+
+using DiffRules: @define_diffrule
+using IrrationalConstants: logtwo, logten
+using NaNMath: NaNMath
+
+###########
+# NaNMath #
+###########
+
+# unary #
+#-------#
+
+@define_diffrule NaNMath.sqrt(x)   = :(  inv(2 * $(NaNMath.sqrt)($x))                     )
+@define_diffrule NaNMath.sin(x)    = :(  $(NaNMath.cos)($x)                               )
+@define_diffrule NaNMath.cos(x)    = :( -$(NaNMath.sin)($x)                               )
+@define_diffrule NaNMath.tan(x)    = :(  1 + $(NaNMath.pow)($(NaNMath.tan)($x), 2)        )
+@define_diffrule NaNMath.asin(x)   = :(  inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2)))  )
+@define_diffrule NaNMath.acos(x)   = :( -inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2)))  )
+@define_diffrule NaNMath.acosh(x)  = :(  inv($(NaNMath.sqrt)($(NaNMath.pow)($x, 2) - 1))  )
+@define_diffrule NaNMath.atanh(x)  = :(  inv(1 - $(NaNMath.pow)($x, 2))                   )
+@define_diffrule NaNMath.log(x)    = :(  inv($x)                                          )
+@define_diffrule NaNMath.log2(x)   = :(  inv($logtwo * $x)                                )
+@define_diffrule NaNMath.log10(x)  = :(  inv($logten * $x)                                )
+@define_diffrule NaNMath.log1p(x)  = :(  inv($x + 1)                                      )
+
+# binary #
+#--------#
+
+@define_diffrule NaNMath.pow(x, y) = :( $y * $(NaNMath.pow)($x, ($y - 1)) ), :( $(NaNMath.pow)($x, $y) * $(NaNMath.log)($x) )
+@define_diffrule NaNMath.max(x, y) = :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
+                                     :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($y), zero($y))))
+@define_diffrule NaNMath.min(x, y) = :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
+                                     :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($x), zero($x))))
+
+end # module

ext/NaNMathSpecialFunctionsExt.jl

@@ -0,0 +1,12 @@
+module NaNMathSpecialFunctionsExt
+
+# `NaNMath.lgamma` is the one rule whose derivative reaches into another package,
+# so it needs both loaded (JuliaDiff/DiffRules.jl#106).
+
+using DiffRules: @define_diffrule
+using NaNMath: NaNMath
+using SpecialFunctions: digamma
+
+@define_diffrule NaNMath.lgamma(x) = :(  $digamma($x)  )
+
+end # module

ext/SpecialFunctionsExt.jl

@@ -0,0 +1,186 @@
+module SpecialFunctionsExt
+
+using DiffRules: @define_diffrule
+using IrrationalConstants: sqrtπ, invsqrtπ
+using SpecialFunctions:
+    SpecialFunctions,
+    airyai,
+    airyaiprime,
+    airyaiprimex,
+    airyaix,
+    airybi,
+    airybiprime,
+    airybiprimex,
+    airybix,
+    besselh,
+    besselhx,
+    besseli,
+    besselix,
+    besselj,
+    besselj0,
+    besselj1,
+    besseljx,
+    besselk,
+    besselkx,
+    bessely,
+    bessely0,
+    bessely1,
+    besselyx,
+    beta,
+    dawson,
+    digamma,
+    ellipe,
+    ellipk,
+    erfcinv,
+    erfcx,
+    erfinv,
+    expint,
+    gamma,
+    hankelh1,
+    hankelh1x,
+    hankelh2,
+    hankelh2x,
+    invdigamma,
+    logerfc,
+    polygamma,
+    trigamma,
+    zeta
+
+####################
+# SpecialFunctions #
+####################
+
+# unary #
+#-------#
+
+@define_diffrule SpecialFunctions.gamma(x)       = :(  $digamma($x) * $gamma($x)          )
+@define_diffrule SpecialFunctions.loggamma(x)    = :(  $digamma($x)                       )
+
+@define_diffrule SpecialFunctions.erf(x)         = :(  2 * ($invsqrtπ * exp(-$x^2))       )
+@define_diffrule SpecialFunctions.erfinv(x)      = :(  ($sqrtπ * exp($erfinv($x)^2)) / 2  )
+@define_diffrule SpecialFunctions.erfc(x)        = :( -($invsqrtπ * exp(-$x^2) * 2)       )
+@define_diffrule SpecialFunctions.logerfc(x)     =
+    :( -2 * ($invsqrtπ * exp(-$x^2 - $logerfc($x)))  )
+
+@define_diffrule SpecialFunctions.erfcinv(x)     = :( -($sqrtπ * exp($erfcinv($x)^2)) / 2 )
+@define_diffrule SpecialFunctions.erfi(x)        = :(  $invsqrtπ * exp($x^2) * 2          )
+@define_diffrule SpecialFunctions.erfcx(x)       = :(  2 * (($x * $erfcx($x)) - $invsqrtπ))
+@define_diffrule SpecialFunctions.logerfcx(x)    = :(  2 * ($x - inv($erfcx($x) * $sqrtπ)))
+
+@define_diffrule SpecialFunctions.dawson(x)      = :(  1 - (2 * $x * $dawson($x))         )
+@define_diffrule SpecialFunctions.digamma(x)     = :(  $trigamma($x)                      )
+@define_diffrule SpecialFunctions.invdigamma(x)  = :(  inv($trigamma($invdigamma($x)))    )
+@define_diffrule SpecialFunctions.trigamma(x)    = :(  $polygamma(2, $x)                  )
+
+# derivatives for `airybix` and `airybiprimex` are only correct for real inputs
+# `airyaix` and `airyaiprimex` are only defined for positive real inputs
+# `airybix` and `airybiprimex` are unscaled for negative real inputs
+@define_diffrule SpecialFunctions.airyai(x)       = :(  $airyaiprime($x)  )
+@define_diffrule SpecialFunctions.airyaiprime(x)  = :(  $x * $airyai($x)  )
+@define_diffrule SpecialFunctions.airyaix(x)      =
+    :(  $airyaiprimex($x) + sqrt($x) * $airyaix($x)  )
+@define_diffrule SpecialFunctions.airyaiprimex(x) =
+    :(  $x * $airyaix($x) + sqrt($x) * $airyaiprimex($x)  )
+@define_diffrule SpecialFunctions.airybi(x)       = :(  $airybiprime($x)  )
+@define_diffrule SpecialFunctions.airybiprime(x)  = :(  $x * $airybi($x)  )
+@define_diffrule SpecialFunctions.airybix(x)      =
+    :(  if $x > zero($x)
+            $airybiprimex($x) - sqrt($x) * $airybix($x)
+        else
+            $airybiprimex($x)
+        end  )
+@define_diffrule SpecialFunctions.airybiprimex(x) =
+    :(  if $x > zero($x)
+            $x * $airybix($x) - sqrt($x) * $airybiprimex($x)
+        else
+            $x * $airybix($x)
+        end  )
+
+@define_diffrule SpecialFunctions.besselj0(x)    = :( -$besselj1($x)                         )
+@define_diffrule SpecialFunctions.besselj1(x)    = :(  ($besselj0($x) - $besselj(2, $x)) / 2 )
+@define_diffrule SpecialFunctions.bessely0(x)    = :( -$bessely1($x)                         )
+@define_diffrule SpecialFunctions.bessely1(x)    = :(  ($bessely0($x) - $bessely(2, $x)) / 2 )
+
+@define_diffrule SpecialFunctions.sinint(x)      = :(  sinc($x / π)  )
+@define_diffrule SpecialFunctions.cosint(x)      = :(  cos($x) / $x  )
+
+@define_diffrule SpecialFunctions.ellipk(m)      =
+    :( ($ellipe($m) / (1 - $m) - $ellipk($m)) / (2 * $m) )
+@define_diffrule SpecialFunctions.ellipe(m)      =
+    :( ($ellipe($m) - $ellipk($m)) / (2 * $m) )
+
+@define_diffrule SpecialFunctions.expint(x)      = :( -exp(-$x) / $x )
+
+# TODO:
+#
+# eta
+# zeta
+
+# binary #
+#--------#
+
+@define_diffrule SpecialFunctions.erf(x, y) =
+    :(  -2 * ($invsqrtπ * exp(-$x^2))  ), :(  2 * ($invsqrtπ * exp(-$y^2))  )
+
+# derivatives with respect to the order `ν` exist but are not implemented
+# (analogously to the ChainRules definitions in SpecialFunctions)
+
+# derivatives for `besselix`, `besseljx` and `besselyx` are only correct for real inputs
+# see https://github.com/JuliaMath/SpecialFunctions.jl/blob/master/src/chainrules.jl
+# for forward-mode and reverse-mode derivatives for complex inputs
+
+@define_diffrule SpecialFunctions.besselj(ν, x)   =
+    :NaN, :(  ($besselj($ν - 1, $x) - $besselj($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besseljx(ν, x)  =
+    :NaN, :(  ($besseljx($ν - 1, $x) - $besseljx($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besseli(ν, x)   =
+    :NaN, :(  ($besseli($ν - 1, $x) + $besseli($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselix(ν, x)  =
+    :NaN, :(  ($besselix($ν - 1, $x) + $besselix($ν + 1, $x)) / 2 - sign($x) * $besselix($ν, $x)  )
+@define_diffrule SpecialFunctions.bessely(ν, x)   =
+    :NaN, :(  ($bessely($ν - 1, $x) - $bessely($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselyx(ν, x)  =
+    :NaN, :(  ($besselyx($ν - 1, $x) - $besselyx($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselk(ν, x)   =
+    :NaN, :( -($besselk($ν - 1, $x) + $besselk($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselkx(ν, x)  =
+    :NaN, :( -($besselkx($ν - 1, $x) + $besselkx($ν + 1, $x)) / 2 + $besselkx($ν, $x)  )
+@define_diffrule SpecialFunctions.besselh(ν, x)   =
+    :NaN, :(  ($besselh($ν - 1, $x) - $besselh($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselhx(ν, x)  =
+    :NaN, :(  ($besselhx($ν - 1, $x) - $besselhx($ν + 1, $x)) / 2 - im * $besselhx($ν, $x)  )
+@define_diffrule SpecialFunctions.hankelh1(ν, x)  =
+    :NaN, :(  ($hankelh1($ν - 1, $x) - $hankelh1($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.hankelh1x(ν, x) =
+    :NaN, :(  ($hankelh1x($ν - 1, $x) - $hankelh1x($ν + 1, $x)) / 2 - im * $hankelh1x($ν, $x)  )
+@define_diffrule SpecialFunctions.hankelh2(ν, x)  =
+    :NaN, :(  ($hankelh2($ν - 1, $x) - $hankelh2($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.hankelh2x(ν, x) =
+    :NaN, :(  ($hankelh2x($ν - 1, $x) - $hankelh2x($ν + 1, $x)) / 2 + im * $hankelh2x($ν, $x)  )
+
+@define_diffrule SpecialFunctions.polygamma(m, x) =
+    :NaN, :(  $polygamma($m + 1, $x)  )
+
+@define_diffrule SpecialFunctions.beta(a, b)      =
+    :(  $beta($a, $b) * ($digamma($a) - $digamma($a + $b))  ),
+    :(  $beta($a, $b) * ($digamma($b) - $digamma($a + $b))  )
+@define_diffrule SpecialFunctions.logbeta(a, b)   =
+    :(  $digamma($a) - $digamma($a + $b)  ), :(  $digamma($b) - $digamma($a + $b)  )
+
+# derivative wrt to `ν` is not implemented
+@define_diffrule SpecialFunctions.expint(ν, x)    =
+    :NaN, :( -$expint($ν - 1, $x) )
+
+# derivative wrt to `s` is not implemented
+@define_diffrule SpecialFunctions.zeta(s, z)      =
+    :NaN, :( -$s * $zeta($s + 1, $z) )
+
+# ternary #
+#---------#
+
+# TODO:
+#
+# besselh
+# besselhx
+
+end # module

src/DiffRules.jl

@@ -2,7 +2,7 @@ __precompile__()
 
 module DiffRules
 
-using IrrationalConstants: logtwo, logten, twoπ, sqrtπ, invsqrtπ
+using IrrationalConstants: logtwo, logten, twoπ
 
 include("api.jl")
 include("rules.jl")