bipass all Diffractor machinery if there is no partials

Author: oxinabox

src/stage1/forward.jl

@@ -137,6 +137,18 @@ function (::∂☆shuffle{N})(args::AbstractTangentBundle{N}...) where {N}
     ∂☆p(ZeroBundle{N-1}(frule), #= ZeroBundle{N-1}(DiffractorRuleConfig()), =# tupargs, map(primal, downargs)...)
 end
 
+# Special shortcut case if there is no derivative information at all:
+function (::∂☆internal{N})(f::AbstractZeroBundle{N}, args::AbstractZeroBundle{N}...) where {N}
+    f_v = primal(f)
+    args_v = map(primal, args)
+    return ZeroBundle{N}(f_v(args_v...))
+end
+function (::∂☆internal{1})(f::AbstractZeroBundle{1}, args::AbstractZeroBundle{1}...)
+    f_v = primal(f)
+    args_v = map(primal, args)
+    return ZeroBundle{1}(f_v(args_v...))
+end
+
 function (::∂☆internal{N})(args::AbstractTangentBundle{N}...) where {N}
     r = ∂☆shuffle{N}()(args...)
     if primal(r) === nothing
@@ -147,6 +159,7 @@ function (::∂☆internal{N})(args::AbstractTangentBundle{N}...) where {N}
 end
 (::∂☆{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆internal{N}()(args...)
 
+
 # Special case rules for performance
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TangentBundle{N}, s::AbstractTangentBundle{N}) where {N}
     s = primal(s)

test/forward.jl

@@ -1,6 +1,6 @@
 module forward_tests
 using Diffractor
-using Diffractor: var"'", ∂⃖, DiffractorRuleConfig
+using Diffractor: var"'", ∂⃖, DiffractorRuleConfig, ZeroBundle
 using ChainRules
 using ChainRulesCore
 using ChainRulesCore: ZeroTangent, NoTangent, frule_via_ad, rrule_via_ad
@@ -50,5 +50,23 @@ let var"'" = Diffractor.PrimeDerivativeFwd
 end
 
 
+@testset "No partials" begin
+    primal_calls = Ref(0)
+    function foo(x, y)
+        primal_calls[]+=1
+        return x+y
+    end
+    
+    frule_calls = Ref(0)
+    function ChainRulesCore.frule((_, ẋ, ẏ), ::typeof(foo), x, y)
+        frule_calls[]+=1
+        return x+y, ẋ+ẏ
+    end
+
+    # Special case if there is no derivative information at all:
+    @test (Diffractor.∂☆{1}())(ZeroBundle{1}(foo), ZeroBundle{1}(2.0), ZeroBundle{1}(3.0)) == ZeroBundle{1}(5.0)
+    @test frule_calls[] == 0
+    @test primal_calls[] == 1
+end
 
 end