Function-based ChainRules rule transferring to Enzyme rules

[ChrisRackauckas]

As mentioned in EnzymeAD/Enzyme.jl#805 (comment)

Define a macro that defines the Enzyme rules for a given function based on the ChainRules. For example, @ruletransfer conv(x)

@ruletransfer conv(x) would define an EnzymeRule that naively calls a ChainRules.jl rule for all of the activity types. This would then serve as a suboptimal fallback that can be applied on a single function basis. The purpose of this macro would be to make it easier to setup the adoption of Zygote throughout the deep learning community by allowing any Zygote rule to trivially port to Enzyme, making Enzyme at least as good as Zygote as quickly as possible. The reason to do it on a function by function basis is that not all rules are necessary, with the main ones being the ones over cudnn kernels. They should of course be replaced with lower level rules that more optimize memory usage as needed, but I think such a simple macro would be a great first step.

This means the rule transfer can be pretty naïve. It could just call the ChainRule and then realout .= chainrulesout , i.e. many ChainRules rules will allocate the output while Enzyme can in theory omit the extra allocation, but for the macro it can ignore such optimizations and just ensure it hits coverage first.

[oxinabox]

I will try and find time to do this, but it will probably be a few weeks at least.

[sharanry]

@ChrisRackauckas To check my understanding of what is needed, we need a macro that can generically define Zygote rules functions like below, right?

using Enzyme
import .EnzymeRules: forward
using ChainRules
function EnzymeRules.forward(func::Const{typeof(sin)}, ::Type{<:Duplicated}, x::Duplicated)
    println("Using custom rule!")
    (ret, dret) = ChainRules.frule((0.0, x.dval), func.val, x.val)
    return Duplicated(ret, dret)
end

x = 1.0
dx = 1.0
@show Enzyme.autodiff(Forward, sin, Duplicated, Duplicated(x, dx))

Using custom rule!
autodiff(Forward, sin, Duplicated, Duplicated(x, dx)) = (0.8414709848078965, 0.5403023058681398)

[ChrisRackauckas]

I think we may not go this route anymore? Talking with Billy more directly, we might just start writing rules for Lux. I'll invite you to the channel.