Code inside function with rrule should not run

[gaurav-arya]

MWE:

using Zygote
import ChainRulesCore.rrule
using ChainRulesCore: NoTangent
using AbstractDifferentiation

function myfunc(x)
    println("This should not print if I have an rrule.")
    x
end

rrule(::typeof(myfunc), x) = (x, (y -> (NoTangent(), y)))

println("Zygote run:")
Zygote.gradient(myfunc, 1) # nothing prints
println("AD run:")
AD.derivative(AD.ZygoteBackend(), myfunc, 1) # something prints

The code inside of myfunc should never run. In addition to possible inefficiency, this may lead to incorrect/confusing results for stateful functions or stochastic functions.

[devmotion]

My hypothesis (without checking it in detail) is that the problem is that the backend for ChainRulesCore integration (used also for Zygote in your example) uses the rrule to define the pullback function (

AbstractDifferentiation.jl/src/ruleconfig.jl

Line 10 in 12a0e27

AD.@primitive function pullback_function(ab::ReverseRuleConfigBackend, f, xs...)

) but does drop the value of the primal completely, which therefore has to be recovered from the original function in value_and_pullback_function (

AbstractDifferentiation.jl/src/AbstractDifferentiation.jl

Line 248 in 12a0e27

function value_and_pullback_function(

) which is e.g. used by derivative (

AbstractDifferentiation.jl/src/AbstractDifferentiation.jl

Line 38 in 12a0e27

function derivative(ab::AbstractBackend, f, xs::Number...)

) via jacobian (

AbstractDifferentiation.jl/src/AbstractDifferentiation.jl

Line 567 in 12a0e27

function AbstractDifferentiation.jacobian($(args...),)

).

I think one "just" has to define value_and_pullback_function in terms of only rrule for these backends. This should avoid the unnecessary call of the original function.

[devmotion]

Just noticed, there's already #34, #35, #36.