SFext

Author: zhujch1

Project.toml

@@ -7,12 +7,16 @@ version = "0.2.1"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesOverloadGeneration = "f51149dc-2911-5acf-81fc-2076a2a81d4f"
-IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+
+[extensions]
+TaylorDiffSFExt = ["SpecialFunctions"]
+
 [compat]
 ChainRules = "1"
 ChainRulesCore = "1"

ext/TaylorDiffSFExt.jl

@@ -0,0 +1,33 @@
+module TaylorDiffSFExt
+using TaylorDiff, SpecialFunctions
+using Symbolics: @variables
+using SymbolicUtils, SymbolicUtils.Code
+using SymbolicUtils: Pow
+using TaylorDiff: value, raise
+using ChainRules, ChainRulesCore
+
+dummy = (NoTangent(), 1)
+@variables z
+for func in (erf, )
+    F = typeof(func)
+    # base case
+    @eval function (op::$F)(t::TaylorScalar{T, 2}) where {T}
+        t0, t1 = value(t)
+        TaylorScalar{T, 2}(frule((NoTangent(), t1), op, t0))
+    end
+    der = frule(dummy, func, z)[2]
+    term, raiser = der isa Pow && der.exp == -1 ? (der.base, raiseinv) : (der, raise)
+    # recursion by raising
+    @eval @generated function (op::$F)(t::TaylorScalar{T, N}) where {T, N}
+        der_expr = $(QuoteNode(toexpr(term)))
+        f = $func
+        quote
+            $(Expr(:meta, :inline))
+            z = TaylorScalar{T, N - 1}(t)
+            df = $der_expr
+            $$raiser($f(value(t)[1]), df, t)
+        end
+    end
+end
+
+end

src/codegen.jl

@@ -1,7 +1,5 @@
 using ChainRules
 using ChainRulesCore
-using SpecialFunctions
-using IrrationalConstants: sqrtπ
 using Symbolics: @variables
 using SymbolicUtils, SymbolicUtils.Code
 using SymbolicUtils: Pow
@@ -13,7 +11,7 @@ for func in (+, -, deg2rad, rad2deg,
     asin, acos, atan, asec, acsc, acot,
     log, log10, log1p, log2,
     asinh, acosh, atanh, asech, acsch,
-    acoth, erf)
+    acoth)
     F = typeof(func)
     # base case
     @eval function (op::$F)(t::TaylorScalar{T, 2}) where {T}