From 5c9332caf9c8de285b20a0f2ea4f730726a4711a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 02:29:04 +0100 Subject: [PATCH 1/2] Take higher-order derivative expressions from DiffRules `get_term_raiser` evaluated ChainRules' `frule` on a Symbolics variable, which breaks as soon as a rule branches on values (#113, ChainRules 1.74.11). DiffRules' expressions are value-free by design, so take the derivative from `DiffRules.diffrule` and drop the Symbolics and SymbolicUtils dependencies. Also import the names used by the code generated in `TaylorDiffSFExt`, which was broken. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 8 ++-- ext/TaylorDiffSFExt.jl | 3 ++ src/utils.jl | 83 ++++++++++++------------------------------ test/Project.toml | 1 + test/primitive.jl | 44 ++++++++++++++++++++++ 5 files changed, 74 insertions(+), 65 deletions(-) diff --git a/Project.toml b/Project.toml index e1ef6b1..63cb96e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,14 +1,13 @@ name = "TaylorDiff" uuid = "b36ab563-344f-407b-a36a-4f200bebf99c" authors = ["Songchen Tan "] -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" diff --git a/ext/TaylorDiffSFExt.jl b/ext/TaylorDiffSFExt.jl index 51d21ba..7ab8831 100644 --- a/ext/TaylorDiffSFExt.jl +++ b/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, raise, raiseinv +using TaylorDiff.ChainRulesCore: frule, NoTangent, zero_tangent for func in (erf, erfc, erfcinv, erfcx, erfi) TaylorDiff.define_unary_function(func, TaylorDiffSFExt) diff --git a/src/utils.jl b/src/utils.jl index 03b75a4..3cca00f 100644 --- a/src/utils.jl +++ b/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`](@ref); otherwise the +derivative is raised with the multiplication rule [`raise`](@ref). - 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)) diff --git a/test/Project.toml b/test/Project.toml index 6610f61..a2b0b86 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/primitive.jl b/test/primitive.jl index cdd38c3..63546c3 100644 --- a/test/primitive.jl +++ b/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 From 7ddcf2d01b76b20f6d8501952bca773f9d1e38fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 13:48:21 +0100 Subject: [PATCH 2/2] Fix formatting, docs cross-references and extension imports - Format the new Float32 testset as JuliaFormatter (sciml) expects - Don't link to the undocumented `raise`/`raiseinv` in the `get_term_raiser` docstring; it is rendered by `@autodocs` and broke the docs build - Don't import `raise`/`raiseinv` in TaylorDiffSFExt; the generated code interpolates the function objects Co-Authored-By: Claude Opus 5.5 --- ext/TaylorDiffSFExt.jl | 2 +- src/utils.jl | 4 ++-- test/primitive.jl | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/TaylorDiffSFExt.jl b/ext/TaylorDiffSFExt.jl index 7ab8831..70a0537 100644 --- a/ext/TaylorDiffSFExt.jl +++ b/ext/TaylorDiffSFExt.jl @@ -1,7 +1,7 @@ module TaylorDiffSFExt using TaylorDiff, SpecialFunctions # Names used by the code that `define_unary_function` generates in this module -using TaylorDiff: TaylorScalar, value, partials, raise, raiseinv +using TaylorDiff: TaylorScalar, value, partials using TaylorDiff.ChainRulesCore: frule, NoTangent, zero_tangent for func in (erf, erfc, erfcinv, erfcx, erfi) diff --git a/src/utils.jl b/src/utils.jl index 3cca00f..9880515 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -12,8 +12,8 @@ using MacroTools: prewalk, postwalk, @capture 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`](@ref); otherwise the -derivative is raised with the multiplication rule [`raise`](@ref). +(respectively `-term`) is raised with the division rule `raiseinv`; otherwise the +derivative is raised with the multiplication rule `raise`. 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, diff --git a/test/primitive.jl b/test/primitive.jl index 63546c3..e622041 100644 --- a/test/primitive.jl +++ b/test/primitive.jl @@ -117,7 +117,7 @@ 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) + atanh, abs, sign, erf, erfc, erfcx, erfi) @test derivative(f, 0.3f0, Val(3)) isa Float32 end end