diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f4887f9..fa06086 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Project.toml b/Project.toml index e30e835..872978d 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/index.md b/docs/src/index.md index 4d21308..faa9b71 100644 --- a/docs/src/index.md +++ b/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 diff --git a/ext/LogExpFunctionsExt.jl b/ext/LogExpFunctionsExt.jl new file mode 100644 index 0000000..71a920b --- /dev/null +++ b/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 diff --git a/ext/NaNMathExt.jl b/ext/NaNMathExt.jl new file mode 100644 index 0000000..c22be76 --- /dev/null +++ b/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 diff --git a/ext/NaNMathSpecialFunctionsExt.jl b/ext/NaNMathSpecialFunctionsExt.jl new file mode 100644 index 0000000..46425cb --- /dev/null +++ b/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 diff --git a/ext/SpecialFunctionsExt.jl b/ext/SpecialFunctionsExt.jl new file mode 100644 index 0000000..34da243 --- /dev/null +++ b/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 diff --git a/src/DiffRules.jl b/src/DiffRules.jl index 67c76f2..0ce877b 100644 --- a/src/DiffRules.jl +++ b/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") diff --git a/src/api.jl b/src/api.jl index a045686..46019d9 100644 --- a/src/api.jl +++ b/src/api.jl @@ -1,13 +1,10 @@ - -const DEFINED_DIFFRULES = Dict{Tuple{Union{Expr,Symbol},Symbol,Int},Any}() - """ @define_diffrule M.f(x) = :(df_dx(\$x)) @define_diffrule M.f(x, y) = :(df_dx(\$x, \$y)), :(df_dy(\$x, \$y)) ⋮ Define a new differentiation rule for the function `M.f` and the given arguments, which should -be treated as bindings to Julia expressions. Return the defined rule's key. +be treated as bindings to Julia expressions. The LHS should be a function call with a non-splatted argument list, and the RHS should be the derivative expression, or in the `n`-ary case, an `n`-tuple of expressions where the @@ -16,84 +13,103 @@ interpolated wherever they are used on the RHS. Note that differentiation rules are purely symbolic, so no type annotations should be used. +Each rule is a method of [`diffrule`](@ref), so rules defined in other packages and in +package extensions are precompiled and visible like any other method. + +A rule is evaluated by whoever asks for it, which need not have the defining package in +scope. Interpolate the functions the RHS calls, rather than naming them, so that the +returned expression stands on its own. + # Examples ```julia -@define_diffrule Base.cos(x) = :(-sin(\$x)) -@define_diffrule Base.:/(x, y) = :(inv(\$y)), :(-\$x / (\$y^2)) -@define_diffrule Base.polygamma(m, x) = :NaN, :(polygamma(\$m + 1, \$x)) +@define_diffrule Base.cos(x) = :(-sin(\$x)) +@define_diffrule Base.:/(x, y) = :(inv(\$y)), :(-\$x / (\$y^2)) +@define_diffrule Base.ldexp(x, y) = :(exp2(\$y)), :NaN +@define_diffrule MyPkg.f(x) = :(\$(MyPkg.g)(\$x)) ``` """ macro define_diffrule(def) @assert isa(def, Expr) && def.head == :(=) "Diff rule expression does not have a left and right side" - lhs = def.args[1] - rhs = def.args[2] + lhs, rhs = def.args @assert isa(lhs, Expr) && lhs.head == :call "LHS is not a function call" - qualified_f = lhs.args[1] - @assert isa(qualified_f, Expr) && qualified_f.head == :(.) "Function is not qualified by module" - M = qualified_f.args[1] - f = _get_quoted_symbol(qualified_f.args[2]) + f = lhs.args[1] + @assert isa(f, Expr) && f.head == :(.) "Function is not qualified by module" args = lhs.args[2:end] - rule = Expr(:->, Expr(:tuple, args...), rhs) - key = Expr(:tuple, Expr(:quote, M), Expr(:quote, f), length(args)) - return esc(quote - $DiffRules.DEFINED_DIFFRULES[$key] = $rule - $key - end) + return esc(:($DiffRules.diffrule(::typeof($f), $(args...)) = $rhs)) end """ - diffrule(M::Union{Expr,Symbol}, f::Symbol, args...) + diffrule(f, args...) + diffrule(M::Union{Module,Symbol}, f::Symbol, args...) -Return the derivative expression for `M.f` at the given argument(s), with the argument(s) +Return the derivative expression for `f` at the given argument(s), with the argument(s) interpolated into the returned expression. In the `n`-ary case, an `n`-tuple of expressions will be returned where the `i`th expression is the derivative of `f` w.r.t the `i`th argument. +Throw a `KeyError` if `M.f` cannot be resolved, which for a package's functions is also the +case as long as that package is not loaded. + # Examples ```jldoctest -julia> DiffRules.diffrule(:Base, :sin, 1) +julia> DiffRules.diffrule(sin, 1) :(cos(1)) -julia> DiffRules.diffrule(:Base, :sin, :x) +julia> DiffRules.diffrule(Base, :sin, :x) :(cos(x)) julia> DiffRules.diffrule(:Base, :sin, :(x * y^2)) :(cos(x * y ^ 2)) ``` """ -diffrule(M::Union{Expr,Symbol}, f::Symbol, args...) = DEFINED_DIFFRULES[M,f,length(args)](args...) +function diffrule end + +function diffrule(M::Module, f::Symbol, args...) + isdefined(M, f) || throw(KeyError((nameof(M), f, length(args)))) + return diffrule(getproperty(M, f), args...) +end + +function diffrule(M::Symbol, f::Symbol, args...) + fn = _resolve(M, f) + fn === nothing && throw(KeyError((M, f, length(args)))) + return diffrule(fn, args...) +end """ - hasdiffrule(M::Union{Expr,Symbol}, f::Symbol, arity::Int) + hasdiffrule(f, arity::Int) + hasdiffrule(M::Union{Module,Symbol}, f::Symbol, arity::Int) -Return `true` if a differentiation rule is defined for `M.f` and `arity`, or return `false` -otherwise. +Return `true` if a differentiation rule is defined for `f` and `arity`, or return `false` +otherwise. Here, `arity` refers to the number of arguments accepted by `f`. -Here, `arity` refers to the number of arguments accepted by `f`. +Rules for a package's functions exist only once that package is loaded, so a query for an +unloaded package returns `false`. # Examples ```jldoctest -julia> DiffRules.hasdiffrule(:Base, :sin, 1) +julia> DiffRules.hasdiffrule(sin, 1) true -julia> DiffRules.hasdiffrule(:Base, :sin, 2) +julia> DiffRules.hasdiffrule(sin, 2) false -julia> DiffRules.hasdiffrule(:Base, :-, 1) -true - julia> DiffRules.hasdiffrule(:Base, :-, 2) true - -julia> DiffRules.hasdiffrule(:Base, :-, 3) -false ``` """ -hasdiffrule(M::Union{Expr,Symbol}, f::Symbol, arity::Int) = haskey(DEFINED_DIFFRULES, (M, f, arity)) +hasdiffrule(f, arity::Int) = hasmethod(diffrule, Tuple{typeof(f),Vararg{Any,arity}}) + +hasdiffrule(M::Module, f::Symbol, arity::Int) = + isdefined(M, f) && hasdiffrule(getproperty(M, f), arity) + +function hasdiffrule(M::Symbol, f::Symbol, arity::Int) + fn = _resolve(M, f) + return fn !== nothing && hasdiffrule(fn, arity) +end # show a deprecation warning if `filter_modules` in `diffrules()` is specified implicitly # we use a custom singleton to figure out if the keyword argument was set explicitly @@ -119,13 +135,14 @@ end diffrules(; filter_modules=(:Base, :SpecialFunctions, :NaNMath)) Return a list of keys that can be used to access all defined differentiation rules for -modules in `filter_modules`. +modules in `filter_modules`. To include all rules, specify `filter_modules = nothing`. -Each key is of the form `(M::Symbol, f::Symbol, arity::Int)`. -Here, `arity` refers to the number of arguments accepted by `f` and `M` is one of the -modules in `filter_modules`. +Each key is of the form `(M::Symbol, f::Symbol, arity::Int)`, where `M` is the name of the +package defining `f` and `arity` is the number of arguments accepted by `f`. -To include all rules, specify `filter_modules = nothing`. +Keys are collected from the method table of [`diffrule`](@ref), so rules defined in other +packages are included. Rules for a package's functions exist only once that package is +loaded: querying before `using SpecialFunctions` will not list its rules. !!! note Calling `diffrules()` with the implicit default keyword argument `filter_modules` @@ -145,55 +162,49 @@ true julia> (:Base, :log, 1) in DiffRules.diffrules() true -julia> (:Base, :*, 2) in DiffRules.diffrules() -true -``` - -If you call `diffrules()`, only rules for Base, SpecialFunctions, and -NaNMath are returned but no rules for LogExpFunctions: -```jldoctest -julia> any(M === :LogExpFunctions for (M, _, _) in DiffRules.diffrules()) -false -``` - -If you set `filter_modules=nothing`, all rules defined in DiffRules are -returned and in particular also rules for LogExpFunctions: -```jldoctest -julia> any( - M === :LogExpFunctions - for (M, _, _) in DiffRules.diffrules(; filter_modules=nothing) - ) -true -``` - -If you set `filter_modules=(:Base,)` only rules for functions in Base are -returned: -```jldoctest julia> all(M === :Base for (M, _, _) in DiffRules.diffrules(; filter_modules=(:Base,))) true ``` """ function diffrules(; filter_modules=DefaultFilterModules()) modules = deprecated_modules(filter_modules) - return if modules === nothing - keys(DEFINED_DIFFRULES) - else - Iterators.filter(keys(DEFINED_DIFFRULES)) do (M, _, _) - return M in modules - end - end + rules = [(_pkgname(fn), nameof(fn), arity) for (fn, arity) in _rules()] + modules === nothing && return rules + return filter(r -> r[1] in modules, rules) end -# For v0.6 and v0.7 compatibility, need to support having the diff rule function enter as a -# `Expr(:quote...)` and a `QuoteNode`. When v0.6 support is dropped, the function will -# always enter in a `QuoteNode` (#23885). -function _get_quoted_symbol(ex::Expr) - @assert ex.head == :quote - @assert length(ex.args) == 1 && isa(ex.args[1], Symbol) "Function not a single symbol" - ex.args[1] +# `parentmodule` reports submodules such as `Base.Math`, which callers cannot splice into +# `M.f`; the root module is the package name they expect. +_pkgname(fn) = nameof(Base.moduleroot(parentmodule(fn))) + +# Rules are the methods of `diffrule` whose first parameter is a singleton function type. +# The `Symbol`/`Module` methods above are not, and are skipped. +function _rule(m::Method) + m.isva && return nothing + params = Base.unwrap_unionall(m.sig).parameters + length(params) >= 2 || return nothing + T = params[2] + T isa DataType && T <: Function && isdefined(T, :instance) || return nothing + return (T.instance, length(params) - 2) end -function _get_quoted_symbol(ex::QuoteNode) - @assert isa(ex.value, Symbol) "Function not a single symbol" - ex.value +function _rules() + rules = Tuple{Function,Int}[] + for m in methods(diffrule) + rule = _rule(m) + rule === nothing || push!(rules, rule) + end + return rules +end + +# Downstream packages resolve every rule by name while precompiling, so scan the method table +# directly instead of building the full list first. +function _resolve(M::Symbol, f::Symbol) + for m in methods(diffrule) + rule = _rule(m) + rule === nothing && continue + fn = rule[1] + nameof(fn) === f && _pkgname(fn) === M && return fn + end + return nothing end diff --git a/src/rules.jl b/src/rules.jl index caa3078..4221866 100644 --- a/src/rules.jl +++ b/src/rules.jl @@ -61,11 +61,6 @@ @define_diffrule Base.deg2rad(x) = :( deg2rad(one($x)) ) @define_diffrule Base.mod2pi(x) = :( isinteger($x / $twoπ) ? oftype(float($x), NaN) : one(float($x)) ) @define_diffrule Base.rad2deg(x) = :( rad2deg(one($x)) ) -@define_diffrule SpecialFunctions.gamma(x) = - :( SpecialFunctions.digamma($x) * SpecialFunctions.gamma($x) ) -@define_diffrule SpecialFunctions.loggamma(x) = - :( SpecialFunctions.digamma($x) ) - @define_diffrule Base.abs(x) = :( $(_abs_deriv)($x) ) # We provide this hook for special number types like `Interval` @@ -104,212 +99,3 @@ _abs_deriv(x) = signbit(x) ? -one(x) : one(x) @define_diffrule Base.ifelse(p, x, y) = false, :($p), :(!$p) =# - -#################### -# SpecialFunctions # -#################### - -# unary # -#-------# - -@define_diffrule SpecialFunctions.erf(x) = :( 2 * ($invsqrtπ * exp(-$x^2)) ) -@define_diffrule SpecialFunctions.erfinv(x) = - :( ($sqrtπ * exp(SpecialFunctions.erfinv($x)^2)) / 2 ) -@define_diffrule SpecialFunctions.erfc(x) = :( -($invsqrtπ * exp(-$x^2) * 2) ) -@define_diffrule SpecialFunctions.logerfc(x) = - :( - 2 * ($invsqrtπ * exp(- $x^2 - SpecialFunctions.logerfc($x))) ) - -@define_diffrule SpecialFunctions.erfcinv(x) = - :( -($sqrtπ * exp(SpecialFunctions.erfcinv($x)^2)) / 2 ) -@define_diffrule SpecialFunctions.erfi(x) = :( $invsqrtπ * exp($x^2) * 2 ) -@define_diffrule SpecialFunctions.erfcx(x) = - :( 2 * (($x * SpecialFunctions.erfcx($x)) - $invsqrtπ) ) -@define_diffrule SpecialFunctions.logerfcx(x) = - :( 2 * ($x - inv(SpecialFunctions.erfcx($x) * $sqrtπ)) ) - -@define_diffrule SpecialFunctions.dawson(x) = - :( 1 - (2 * $x * SpecialFunctions.dawson($x)) ) -@define_diffrule SpecialFunctions.digamma(x) = - :( SpecialFunctions.trigamma($x) ) -@define_diffrule SpecialFunctions.invdigamma(x) = - :( inv(SpecialFunctions.trigamma(SpecialFunctions.invdigamma($x))) ) -@define_diffrule SpecialFunctions.trigamma(x) = - :( SpecialFunctions.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) = - :( SpecialFunctions.airyaiprime($x) ) -@define_diffrule SpecialFunctions.airyaiprime(x) = - :( $x * SpecialFunctions.airyai($x) ) -@define_diffrule SpecialFunctions.airyaix(x) = - :( SpecialFunctions.airyaiprimex($x) + sqrt($x) * SpecialFunctions.airyaix($x) ) -@define_diffrule SpecialFunctions.airyaiprimex(x) = - :( $x * SpecialFunctions.airyaix($x) + sqrt($x) * SpecialFunctions.airyaiprimex($x) ) -@define_diffrule SpecialFunctions.airybi(x) = - :( SpecialFunctions.airybiprime($x) ) -@define_diffrule SpecialFunctions.airybiprime(x) = - :( $x * SpecialFunctions.airybi($x) ) -@define_diffrule SpecialFunctions.airybix(x) = - :( if $x > zero($x) - SpecialFunctions.airybiprimex($x) - sqrt($x) * SpecialFunctions.airybix($x) - else - SpecialFunctions.airybiprimex($x) - end ) -@define_diffrule SpecialFunctions.airybiprimex(x) = - :( if $x > zero($x) - $x * SpecialFunctions.airybix($x) - sqrt($x) * SpecialFunctions.airybiprimex($x) - else - $x * SpecialFunctions.airybix($x) - end ) - -@define_diffrule SpecialFunctions.besselj0(x) = - :( -SpecialFunctions.besselj1($x) ) -@define_diffrule SpecialFunctions.besselj1(x) = - :( (SpecialFunctions.besselj0($x) - SpecialFunctions.besselj(2, $x)) / 2 ) -@define_diffrule SpecialFunctions.bessely0(x) = - :( -SpecialFunctions.bessely1($x) ) -@define_diffrule SpecialFunctions.bessely1(x) = - :( (SpecialFunctions.bessely0($x) - SpecialFunctions.bessely(2, $x)) / 2 ) - -@define_diffrule SpecialFunctions.sinint(x) = :( sinc($x / π) ) -@define_diffrule SpecialFunctions.cosint(x) = :( cos($x) / $x ) - -@define_diffrule SpecialFunctions.ellipk(m) = - :( (SpecialFunctions.ellipe($m) / (1 - $m) - SpecialFunctions.ellipk($m)) / (2 * $m) ) -@define_diffrule SpecialFunctions.ellipe(m) = - :( (SpecialFunctions.ellipe($m) - SpecialFunctions.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, :( (SpecialFunctions.besselj($ν - 1, $x) - SpecialFunctions.besselj($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besseljx(ν, x) = - :NaN, :( (SpecialFunctions.besseljx($ν - 1, $x) - SpecialFunctions.besseljx($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besseli(ν, x) = - :NaN, :( (SpecialFunctions.besseli($ν - 1, $x) + SpecialFunctions.besseli($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besselix(ν, x) = - :NaN, :( (SpecialFunctions.besselix($ν - 1, $x) + SpecialFunctions.besselix($ν + 1, $x)) / 2 - sign($x) * SpecialFunctions.besselix($ν, $x) ) -@define_diffrule SpecialFunctions.bessely(ν, x) = - :NaN, :( (SpecialFunctions.bessely($ν - 1, $x) - SpecialFunctions.bessely($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besselyx(ν, x) = - :NaN, :( (SpecialFunctions.besselyx($ν - 1, $x) - SpecialFunctions.besselyx($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besselk(ν, x) = - :NaN, :( -(SpecialFunctions.besselk($ν - 1, $x) + SpecialFunctions.besselk($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besselkx(ν, x) = - :NaN, :( -(SpecialFunctions.besselkx($ν - 1, $x) + SpecialFunctions.besselkx($ν + 1, $x)) / 2 + SpecialFunctions.besselkx($ν, $x) ) -@define_diffrule SpecialFunctions.besselh(ν, x) = - :NaN, :( (SpecialFunctions.besselh($ν - 1, $x) - SpecialFunctions.besselh($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.besselhx(ν, x) = - :NaN, :( (SpecialFunctions.besselhx($ν - 1, $x) - SpecialFunctions.besselhx($ν + 1, $x)) / 2 - im * SpecialFunctions.besselhx($ν, $x) ) -@define_diffrule SpecialFunctions.hankelh1(ν, x) = - :NaN, :( (SpecialFunctions.hankelh1($ν - 1, $x) - SpecialFunctions.hankelh1($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.hankelh1x(ν, x) = - :NaN, :( (SpecialFunctions.hankelh1x($ν - 1, $x) - SpecialFunctions.hankelh1x($ν + 1, $x)) / 2 - im * SpecialFunctions.hankelh1x($ν, $x) ) -@define_diffrule SpecialFunctions.hankelh2(ν, x) = - :NaN, :( (SpecialFunctions.hankelh2($ν - 1, $x) - SpecialFunctions.hankelh2($ν + 1, $x)) / 2 ) -@define_diffrule SpecialFunctions.hankelh2x(ν, x) = - :NaN, :( (SpecialFunctions.hankelh2x($ν - 1, $x) - SpecialFunctions.hankelh2x($ν + 1, $x)) / 2 + im * SpecialFunctions.hankelh2x($ν, $x) ) - -@define_diffrule SpecialFunctions.polygamma(m, x) = - :NaN, :( SpecialFunctions.polygamma($m + 1, $x) ) - -@define_diffrule SpecialFunctions.beta(a, b) = - :( SpecialFunctions.beta($a, $b)*(SpecialFunctions.digamma($a) - SpecialFunctions.digamma($a + $b)) ), :( SpecialFunctions.beta($a, $b)*(SpecialFunctions.digamma($b) - SpecialFunctions.digamma($a + $b)) ) -@define_diffrule SpecialFunctions.logbeta(a, b) = - :( SpecialFunctions.digamma($a) - SpecialFunctions.digamma($a + $b) ), :( SpecialFunctions.digamma($b) - SpecialFunctions.digamma($a + $b) ) - -# derivative wrt to `ν` is not implemented -@define_diffrule SpecialFunctions.expint(ν, x) = - :NaN, :( -SpecialFunctions.expint($ν - 1, $x) ) - -# derivative wrt to `s` is not implemented -@define_diffrule SpecialFunctions.zeta(s, z) = - :NaN, :( - $s * SpecialFunctions.zeta($s + 1, $z) ) - -# ternary # -#---------# - -# TODO: -# -# besselh -# besselhx - -########### -# 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) ) -@define_diffrule NaNMath.lgamma(x) = :( SpecialFunctions.digamma($x) ) - - -# 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)))) - -################### -# LogExpFunctions # -################### - -# unary -@define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x)) -@define_diffrule LogExpFunctions.logistic(x) = :(z = LogExpFunctions.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) = :(LogExpFunctions.logistic($x)) -@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - LogExpFunctions.log1mexp($x))) -@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - LogExpFunctions.log2mexp($x))) -@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - LogExpFunctions.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 - LogExpFunctions.logaddexp($x, $y))), :(exp($y - LogExpFunctions.logaddexp($x, $y))) -@define_diffrule LogExpFunctions.logsubexp(x, y) = - :(z = LogExpFunctions.logsubexp($x, $y); $x > $y ? exp($x - z) : -exp($x - z)), - :(z = LogExpFunctions.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) diff --git a/test/baseline.jl b/test/baseline.jl new file mode 100644 index 0000000..8549d49 --- /dev/null +++ b/test/baseline.jl @@ -0,0 +1,159 @@ +# Every differentiation rule DiffRules defines, as (module, function, arity). +# Downstream packages generate code by iterating `diffrules()` and splicing `M.f`, +# so this list guards the keys they rely on against unintended changes. + +const BASELINE = [ + (:Base, :*, 2), + (:Base, :+, 1), + (:Base, :+, 2), + (:Base, :-, 1), + (:Base, :-, 2), + (:Base, :/, 2), + (:Base, :\, 2), + (:Base, :^, 2), + (:Base, :abs, 1), + (:Base, :abs2, 1), + (:Base, :acos, 1), + (:Base, :acosd, 1), + (:Base, :acosh, 1), + (:Base, :acot, 1), + (:Base, :acotd, 1), + (:Base, :acoth, 1), + (:Base, :acsc, 1), + (:Base, :acscd, 1), + (:Base, :acsch, 1), + (:Base, :asec, 1), + (:Base, :asecd, 1), + (:Base, :asech, 1), + (:Base, :asin, 1), + (:Base, :asind, 1), + (:Base, :asinh, 1), + (:Base, :atan, 1), + (:Base, :atan, 2), + (:Base, :atand, 1), + (:Base, :atanh, 1), + (:Base, :cbrt, 1), + (:Base, :cos, 1), + (:Base, :cosd, 1), + (:Base, :cosh, 1), + (:Base, :cospi, 1), + (:Base, :cot, 1), + (:Base, :cotd, 1), + (:Base, :coth, 1), + (:Base, :csc, 1), + (:Base, :cscd, 1), + (:Base, :csch, 1), + (:Base, :deg2rad, 1), + (:Base, :exp, 1), + (:Base, :exp10, 1), + (:Base, :exp2, 1), + (:Base, :expm1, 1), + (:Base, :hypot, 2), + (:Base, :inv, 1), + (:Base, :ldexp, 2), + (:Base, :log, 1), + (:Base, :log, 2), + (:Base, :log10, 1), + (:Base, :log1p, 1), + (:Base, :log2, 1), + (:Base, :max, 2), + (:Base, :min, 2), + (:Base, :mod, 2), + (:Base, :mod2pi, 1), + (:Base, :rad2deg, 1), + (:Base, :rem, 2), + (:Base, :rem2pi, 2), + (:Base, :sec, 1), + (:Base, :secd, 1), + (:Base, :sech, 1), + (:Base, :sin, 1), + (:Base, :sinc, 1), + (:Base, :sind, 1), + (:Base, :sinh, 1), + (:Base, :sinpi, 1), + (:Base, :sqrt, 1), + (:Base, :tan, 1), + (:Base, :tand, 1), + (:Base, :tanh, 1), + (:LogExpFunctions, :log1mexp, 1), + (:LogExpFunctions, :log1pexp, 1), + (:LogExpFunctions, :log1pmx, 1), + (:LogExpFunctions, :log1psq, 1), + (:LogExpFunctions, :log2mexp, 1), + (:LogExpFunctions, :logaddexp, 2), + (:LogExpFunctions, :logexpm1, 1), + (:LogExpFunctions, :logistic, 1), + (:LogExpFunctions, :logit, 1), + (:LogExpFunctions, :logmxp1, 1), + (:LogExpFunctions, :logsubexp, 2), + (:LogExpFunctions, :xlog1py, 2), + (:LogExpFunctions, :xlogx, 1), + (:LogExpFunctions, :xlogy, 2), + (:NaNMath, :acos, 1), + (:NaNMath, :acosh, 1), + (:NaNMath, :asin, 1), + (:NaNMath, :atanh, 1), + (:NaNMath, :cos, 1), + (:NaNMath, :lgamma, 1), + (:NaNMath, :log, 1), + (:NaNMath, :log10, 1), + (:NaNMath, :log1p, 1), + (:NaNMath, :log2, 1), + (:NaNMath, :max, 2), + (:NaNMath, :min, 2), + (:NaNMath, :pow, 2), + (:NaNMath, :sin, 1), + (:NaNMath, :sqrt, 1), + (:NaNMath, :tan, 1), + (:SpecialFunctions, :airyai, 1), + (:SpecialFunctions, :airyaiprime, 1), + (:SpecialFunctions, :airyaiprimex, 1), + (:SpecialFunctions, :airyaix, 1), + (:SpecialFunctions, :airybi, 1), + (:SpecialFunctions, :airybiprime, 1), + (:SpecialFunctions, :airybiprimex, 1), + (:SpecialFunctions, :airybix, 1), + (:SpecialFunctions, :besselh, 2), + (:SpecialFunctions, :besselhx, 2), + (:SpecialFunctions, :besseli, 2), + (:SpecialFunctions, :besselix, 2), + (:SpecialFunctions, :besselj, 2), + (:SpecialFunctions, :besselj0, 1), + (:SpecialFunctions, :besselj1, 1), + (:SpecialFunctions, :besseljx, 2), + (:SpecialFunctions, :besselk, 2), + (:SpecialFunctions, :besselkx, 2), + (:SpecialFunctions, :bessely, 2), + (:SpecialFunctions, :bessely0, 1), + (:SpecialFunctions, :bessely1, 1), + (:SpecialFunctions, :besselyx, 2), + (:SpecialFunctions, :beta, 2), + (:SpecialFunctions, :cosint, 1), + (:SpecialFunctions, :dawson, 1), + (:SpecialFunctions, :digamma, 1), + (:SpecialFunctions, :ellipe, 1), + (:SpecialFunctions, :ellipk, 1), + (:SpecialFunctions, :erf, 1), + (:SpecialFunctions, :erf, 2), + (:SpecialFunctions, :erfc, 1), + (:SpecialFunctions, :erfcinv, 1), + (:SpecialFunctions, :erfcx, 1), + (:SpecialFunctions, :erfi, 1), + (:SpecialFunctions, :erfinv, 1), + (:SpecialFunctions, :expint, 1), + (:SpecialFunctions, :expint, 2), + (:SpecialFunctions, :gamma, 1), + (:SpecialFunctions, :hankelh1, 2), + (:SpecialFunctions, :hankelh1x, 2), + (:SpecialFunctions, :hankelh2, 2), + (:SpecialFunctions, :hankelh2x, 2), + (:SpecialFunctions, :invdigamma, 1), + (:SpecialFunctions, :logbeta, 2), + (:SpecialFunctions, :logerfc, 1), + (:SpecialFunctions, :logerfcx, 1), + (:SpecialFunctions, :loggamma, 1), + (:SpecialFunctions, :polygamma, 2), + (:SpecialFunctions, :sinint, 1), + (:SpecialFunctions, :trigamma, 1), + (:SpecialFunctions, :zeta, 2), +] diff --git a/test/precompilation.jl b/test/precompilation.jl new file mode 100644 index 0000000..8f2aa18 --- /dev/null +++ b/test/precompilation.jl @@ -0,0 +1,120 @@ +# Which rules exist depends on what is loaded, which the test suite itself cannot vary: it loads +# all providers. So probe throwaway environments from fresh Julia processes. + +const LOAD_PATH_SEP = Sys.iswindows() ? ';' : ':' + +# `Pkg.test` runs us with `JULIA_LOAD_PATH="@:/test"`. A child inherits that, losing +# `@stdlib` and gaining our test environment, so pin it to its own project plus the stdlibs. +function probe(dir, script) + cmd = `$(Base.julia_cmd()) --startup-file=no --project=$dir -e $script` + return read( + addenv( + cmd, + "JULIA_LOAD_PATH" => "@$(LOAD_PATH_SEP)@stdlib", + "JULIA_PROJECT" => nothing, + ), + String, + ) +end + +const DIFFRULES_UUID = "b552c78f-8df3-52c6-915a-8e097449b14b" +const NANMATH_UUID = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +const SPECIALFUNCTIONS_UUID = "276daf66-3868-5448-9aa4-cd146d93841b" +const PROBE_UUID = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01" + +function write_env(dir, deps; probe_pkg=false) + mkpath(dir) + entries = join(("$name = \"$uuid\"" for (name, uuid) in deps), "\n") + sources = "DiffRules = { path = \"$(escape_string(pkgdir(DiffRules)))\" }" + if probe_pkg + entries *= "\nProbe = \"$PROBE_UUID\"" + sources *= "\nProbe = { path = \"$(escape_string(joinpath(dir, "Probe")))\" }" + end + write(joinpath(dir, "Project.toml"), """ + [deps] + $entries + + [sources] + $sources + """) + return dir +end + +@testset "precompilation" begin + mktempdir() do root + # `Probe.RULES` records what was visible at `Probe`'s own precompile time. + mkpath(joinpath(root, "both", "Probe", "src")) + write(joinpath(root, "both", "Probe", "Project.toml"), """ + name = "Probe" + uuid = "$PROBE_UUID" + version = "0.1.0" + + [deps] + DiffRules = "$DIFFRULES_UUID" + SpecialFunctions = "$SPECIALFUNCTIONS_UUID" + """) + write(joinpath(root, "both", "Probe", "src", "Probe.jl"), """ + module Probe + using DiffRules, SpecialFunctions + cube(x) = x^3 + DiffRules.@define_diffrule Probe.cube(x) = :(3 * \$x^2) + const RULES = DiffRules.diffrules(; filter_modules=nothing) + end + """) + + both = write_env( + joinpath(root, "both"), + ("DiffRules" => DIFFRULES_UUID, + "NaNMath" => NANMATH_UUID, + "SpecialFunctions" => SPECIALFUNCTIONS_UUID); + probe_pkg=true, + ) + bare = write_env(joinpath(root, "bare"), ("DiffRules" => DIFFRULES_UUID,)) + nanmath = write_env( + joinpath(root, "nanmath"), + ("DiffRules" => DIFFRULES_UUID, "NaNMath" => NANMATH_UUID), + ) + + @testset "rules of a dependent package survive its precompilation" begin + out = probe(both, """ + using Pkg; Pkg.instantiate(; io=devnull) + using Probe, DiffRules, NaNMath + println("rule_survives=", + (:Probe, :cube, 1) in DiffRules.diffrules(; filter_modules=nothing)) + println("rule_resolves=", DiffRules.diffrule(:Probe, :cube, :x) == :(3 * x^2)) + println("ext_at_precompile=", any(r -> r[1] === :SpecialFunctions, Probe.RULES)) + println("lgamma=", DiffRules.hasdiffrule(:NaNMath, :lgamma, 1)) + """) + @test occursin("rule_survives=true", out) + @test occursin("rule_resolves=true", out) + @test occursin("ext_at_precompile=true", out) + @test occursin("lgamma=true", out) + end + + @testset "no provider loaded, no provider rules" begin + out = probe(bare, """ + using Pkg; Pkg.instantiate(; io=devnull) + using DiffRules + println("erf=", DiffRules.hasdiffrule(:SpecialFunctions, :erf, 1)) + println("nanmath_sin=", DiffRules.hasdiffrule(:NaNMath, :sin, 1)) + println("base_only=", all(M === :Base for (M, _, _) in + DiffRules.diffrules(; filter_modules=nothing))) + """) + @test occursin("erf=false", out) + @test occursin("nanmath_sin=false", out) + @test occursin("base_only=true", out) + end + + # `NaNMath.lgamma` differentiates to `SpecialFunctions.digamma` (#106). + @testset "NaNMath alone does not define the lgamma rule" begin + out = probe(nanmath, """ + using Pkg; Pkg.instantiate(; io=devnull) + using DiffRules, NaNMath + println("nanmath_sin=", DiffRules.hasdiffrule(:NaNMath, :sin, 1)) + println("lgamma=", DiffRules.hasdiffrule(:NaNMath, :lgamma, 1)) + """) + @test occursin("nanmath_sin=true", out) + @test occursin("lgamma=false", out) + end + end +end diff --git a/test/registry.jl b/test/registry.jl new file mode 100644 index 0000000..18a8a31 --- /dev/null +++ b/test/registry.jl @@ -0,0 +1,66 @@ +@testset "registry" begin + @testset "baseline" begin + rules = DiffRules.diffrules(; filter_modules=nothing) + @test allunique(rules) + sort!(rules; by = r -> (string(r[1]), string(r[2]), r[3])) + @test rules == BASELINE + @testset "$M.$f/$n" for (M, f, n) in BASELINE + @test DiffRules.hasdiffrule(M, f, n) + end + end + + # `parentmodule` reports `Base.Math` for these, which would break `@eval $M.$f`. + @testset "Base.Math functions report :Base" begin + @testset "$f" for f in (:sec, :sind, :hypot, :cot, :deg2rad, :mod2pi, :rem2pi) + @test any(r -> r[1] === :Base && r[2] === f, + DiffRules.diffrules(; filter_modules=nothing)) + end + end + + @testset "lookup by function" begin + @test DiffRules.diffrule(Base.sin, :x) == :(cos(x)) + @test DiffRules.diffrule(Base.sec, :x) == :(sec(x) * tan(x)) + @test DiffRules.hasdiffrule(Base.sin, 1) + @test !DiffRules.hasdiffrule(Base.sin, 2) + # `log` has a rule at both arities + @test DiffRules.hasdiffrule(Base.log, 1) && DiffRules.hasdiffrule(Base.log, 2) + end + + @testset "lookup by module" begin + @test DiffRules.diffrule(Base, :sin, :x) == :(cos(x)) + @test DiffRules.diffrule(SpecialFunctions, :digamma, :x) == + :($(SpecialFunctions.trigamma)(x)) + @test DiffRules.hasdiffrule(Base, :sin, 1) + @test !DiffRules.hasdiffrule(Base, :sin, 2) + end + + # Rules splice in the functions they call, so evaluating one does not require the + # defining package to be in scope. + @testset "rules evaluate without the package in scope" begin + mod = Module(:NoImports) + d = DiffRules.diffrule(:SpecialFunctions, :digamma, :x) + @test Base.eval(mod, :(let x = 1.5; $d end)) ≈ SpecialFunctions.trigamma(1.5) + end + + @testset "unknown rules" begin + @test !DiffRules.hasdiffrule(:Base, :nonexistent, 1) + @test !DiffRules.hasdiffrule(:NotLoaded, :f, 1) + @test !DiffRules.hasdiffrule(Base, :nonexistent, 1) + @test !DiffRules.hasdiffrule(:Base, :sum, 1) + @test_throws KeyError DiffRules.diffrule(:Base, :nonexistent, :x) + @test_throws KeyError DiffRules.diffrule(Base, :nonexistent, :x) + @test_throws MethodError DiffRules.diffrule(Base.sin, :x, :y) + end + + @testset "filter_modules" begin + @test all(M === :Base for (M, _, _) in DiffRules.diffrules(; filter_modules=(:Base,))) + @test all(M in (:Base, :SpecialFunctions, :NaNMath) + for (M, _, _) in DiffRules.diffrules()) + end + + # undocumented, but overloaded by IntervalArithmetic and Zygote + @testset "_abs_deriv hook" begin + @test DiffRules._abs_deriv(-2.0) == -1.0 + @test DiffRules.diffrule(Base.abs, :x) == :($(DiffRules._abs_deriv)(x)) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index f1fa041..d493b57 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,7 +11,12 @@ Random.seed!(1) # Set `max_range` to avoid domain errors. const finitediff = central_fdm(5, 1, max_range=1e-3) +include("baseline.jl") + @testset "DiffRules" begin +include("registry.jl") +include("precompilation.jl") + @testset "check rules" begin non_diffeable_arg_functions = [(:Base, :rem2pi, 2), (:Base, :ldexp, 2), (:Base, :ifelse, 3)]