Take higher-order derivative expressions from DiffRules instead of evaluating ChainRules rules symbolically

Project.toml

@@ -1,14 +1,13 @@
 name = "TaylorDiff"
 uuid = "b36ab563-344f-407b-a36a-4f200bebf99c"
 authors = ["Songchen Tan <i@tansongchen.com>"]
-version = "0.3.5"
+version = "0.3.6"
 
 [deps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
-Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [weakdeps]
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
@@ -23,10 +22,9 @@ TaylorDiffZygoteExt = ["Zygote"]
 [compat]
 ChainRules = "1"
 ChainRulesCore = "1"
+DiffRules = "1.17"
 MacroTools = "0.5"
 NNlib = "0.9"
 SpecialFunctions = "2"
-SymbolicUtils = "3, 4"
-Symbolics = "6, 7"
 Zygote = "0.6, 0.7"
 julia = "1.10"

ext/TaylorDiffSFExt.jl

@@ -1,5 +1,8 @@
 module TaylorDiffSFExt
 using TaylorDiff, SpecialFunctions
+# Names used by the code that `define_unary_function` generates in this module
+using TaylorDiff: TaylorScalar, value, partials
+using TaylorDiff.ChainRulesCore: frule, NoTangent, zero_tangent
 
 for func in (erf, erfc, erfcinv, erfcx, erfi)
     TaylorDiff.define_unary_function(func, TaylorDiffSFExt)

src/utils.jl

@@ -3,71 +3,34 @@
 
 using ChainRules
 using ChainRulesCore
-using Symbolics: Symbolics, @variables, @rule, unwrap, isdiv
-import SymbolicUtils
-using SymbolicUtils.Code: toexpr
+import DiffRules
 using MacroTools
-using MacroTools: prewalk, postwalk
+using MacroTools: prewalk, postwalk, @capture
 
 """
-Pick a strategy for raising the derivative of a function.
-If the derivative is like 1 over something, raise with the division rule;
-otherwise, raise with the multiplication rule.
-"""
-function get_term_raiser end
-
-@static if pkgversion(Symbolics) < v"7"
-  function get_term_raiser(func)
-      @variables z
-      r1 = @rule -1 * (1 / ~x) => (-1) / ~x
-      der = frule((NoTangent(), true), func, z)[2]
-      term = unwrap(der)
-      maybe_rewrite = r1(term)
-      if maybe_rewrite !== nothing
-          term = maybe_rewrite
-      end
-      if isdiv(term) && (term.num == 1 || term.num == -1)
-          term.den * term.num, raiseinv
-      else
-          term, raise
-      end
-  end
-else
-    const COMMON_Z = only(@variables z)
-    const FALLBACK_RULE = (@rule -1 * (1 / ~x) => (-1) / ~x)
+    get_term_raiser(func)
 
-    function is_plusminus_one(@nospecialize(x))
-        if x isa Int
-            return x == 1 || x == -1
-        elseif x isa Int32
-            return x == 1 || x == -1
-        elseif x isa Float64
-            return x == 1 || x == -1
-        elseif x isa Float32
-            return x == 1 || x == -1
-        elseif x isa Number
-            return (x == 1)::Bool || (x == -1)::Bool
-        else
-            return false
-        end
-    end
+Return the derivative of `func` as an expression in `z`, together with the strategy for raising
+it to higher orders: if the derivative is `inv(term)` or `1 / term` (possibly negated), `term`
+(respectively `-term`) is raised with the division rule `raiseinv`; otherwise the
+derivative is raised with the multiplication rule `raise`.
 
-    function get_term_raiser(func)
-        der = frule((NoTangent(), true), func, COMMON_Z)[2]
-        term = unwrap(der)
-        maybe_rewrite = FALLBACK_RULE(term)
-        if maybe_rewrite !== nothing
-            term = maybe_rewrite
-        end
-        if isdiv(term)
-            num, den = SymbolicUtils.arguments(term)
-            if SymbolicUtils.isconst(num) && is_plusminus_one(SymbolicUtils.unwrap_const(num))
-                return num * den, raiseinv
-            end
-        end
-        return term, raise
-    end
+The expression is evaluated on Taylor series, so it must equal the derivative as a function
+(not only numerically) and must not branch on values. Such expressions are taken from DiffRules,
+whose rules are value-free expressions by design (unlike ChainRules' rules, which are only
+meant to be evaluated on numbers).
+"""
+function get_term_raiser(func)
+    f = nameof(func)
+    M = DiffRules.hasdiffrule(:Base, f, 1) ? :Base : nameof(parentmodule(func))
+    der = DiffRules.diffrule(M, f, :z)
+    @capture(der, inv(den_) | 1 / den_) && return den, raiseinv
+    @capture(der, -inv(den_) | -(1 / den_) | -1 / den_) && return :(-$den), raiseinv
+    return der, raise
 end
+# DiffRules' rule for `abs` branches on the sign of its argument, and it has no rule for `sign`
+get_term_raiser(::typeof(abs)) = :(sign(z)), raise
+get_term_raiser(::typeof(sign)) = :(zero(z)), raise
 
 function define_unary_function(func, m)
     F = typeof(func)
@@ -81,7 +44,7 @@ function define_unary_function(func, m)
     term, raiser = get_term_raiser(func)
     # Higher order: recursion by raising
     @eval m @generated function (op::$F)(t::TaylorScalar{T, p}) where {T, p}
-        expr = $(QuoteNode(toexpr(term)))
+        expr = $(QuoteNode(term))
         f = $func
         quote
             $(Expr(:meta, :inline))

test/Project.toml

@@ -5,6 +5,7 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

test/primitive.jl

@@ -1,4 +1,6 @@
 using FiniteDifferences
+using SpecialFunctions
+using TaylorDiff.ChainRulesCore: frule, NoTangent, AbstractZero
 
 @testset "No derivative or linear" begin
     some_number, another_number = 1.9, 2.6
@@ -77,3 +79,45 @@ end
         @test f(t) == t
     end
 end
+
+@testset "Higher-order terms match the first-order rules" begin
+    # `get_term_raiser` takes the derivative expressions from DiffRules; check them against
+    # ChainRules' rules, which are used for the first order
+    domain = Dict(
+        asec => (1.3, -2.0), acsc => (1.3, -2.0), acosh => (1.3, 2.0), acoth => (1.3, -2.0),
+        log => (0.3, 2.0), log10 => (0.3, 2.0), log2 => (0.3, 2.0), asech => (0.3, 0.7),
+        erfcinv => (0.3, 1.5))
+    for f in (deg2rad, rad2deg, sinh, cosh, tanh, asin, acos, atan, asec, acsc, acot,
+            log, log10, log1p, log2, asinh, acosh, atanh, asech, acsch, acoth, abs, sign,
+            erf, erfc, erfcinv, erfcx, erfi),
+        x in get(domain, f, (0.3, -0.7))
+
+        term, raiser = TaylorDiff.get_term_raiser(f)
+        d = @eval let z = $x
+            $term
+        end
+        expected = frule((NoTangent(), 1.0), f, x)[2]
+        expected isa AbstractZero && (expected = 0.0)  # e.g., `sign`
+        @test (raiser === TaylorDiff.raiseinv ? inv(d) : d) ≈ expected
+    end
+end
+
+@testset "Derivatives of the form 1 / g are raised by division" begin
+    # Guards against changes in the form of the DiffRules expressions that would silently
+    # fall back to the slower and less accurate multiplication rule
+    for f in (asin, acos, atan, asec, acsc, acot, log, log10, log1p, log2,
+        asinh, acosh, atanh, asech, acsch, acoth)
+        @test TaylorDiff.get_term_raiser(f)[2] === TaylorDiff.raiseinv
+    end
+    for f in (deg2rad, rad2deg, sinh, cosh, tanh, abs, sign,
+        erf, erfc, erfcinv, erfcx, erfi)
+        @test TaylorDiff.get_term_raiser(f)[2] === TaylorDiff.raise
+    end
+end
+
+@testset "Higher-order derivatives preserve Float32" begin
+    for f in (deg2rad, rad2deg, sinh, cosh, tanh, atan, acot, log, log10, log1p, log2,
+        atanh, abs, sign, erf, erfc, erfcx, erfi)
+        @test derivative(f, 0.3f0, Val(3)) isa Float32
+    end
+end