New autodiff rule from `DifferentiateWith`

[yebai]

As explored in discussions (e.g., #68), we can introduce a flexible "rule switch" mechanism into DifferentiationInterface.jl (DI) by generalising the existing DifferentiateWith functionality.

The central concept is a macro that allows for the dynamic generation of differentiation rules:

Targtypes = (typeof(x), typeof(y), Constant{typeof(z)}, Cache{typeof(c)})
Tfunc_sig =Tuple{typeof(f), Targtypes...}
@rule_from_DifferentiateWith(AutoTargetAD, ::Tfunc_sig, backend_to_use::AbstractADType)

In this definition:

• @rule_from_DifferentiateWith is a macro designed to generate new differentiation rules.
• AutoTargetAD specifies the automatic differentiation (AD) package for which the new rule is being created (e.g., AutoChainRules).
• func_sig represents the signature of the function f along with its argument types (argtypes...) for which the rule is being defined.
• backend_to_use indicates another AD package (the "backend") whose capabilities will be leveraged to define the rule for AutoTargetAD.

This approach provides a powerful and general way to establish a "rule switch mechanism." It avoids imposing strict assumptions on the rule interface, unlike systems such as ChainRules, offering greater flexibility.

After #806, we could also do

dw = DifferentiateWith(Tfunc_sig, backend_to_use::AbstractADType)
@rule_from_DifferentiateWith(AutoTargetAD, dw)

[gdalle]

I'm not sure I understand where you're going with the last part. To me, the whole purpose of having a macro generate rules is to dispense the user from creating a DifferentiateWith object, by making the naked function f differentiable directly?

Essentially, what I envisioned was a macro

@rule_from_DifferentiateWith(target_backend, DifferentiateWith(f, source_backend))

which would generate code like this when target_backend isa AutoChainRules for example

function ChainRulesCore.rrule(::typeof(f), x)
    dw = DifferentiateWith(f, source_backend)
    return rrule(dw, x)  # already implemented
end

[yebai]

I'm not sure I understand where you're going with the last part.

It is a syntax sugar that allows us to reuse a pre-defined diffwith instance. For example

Tfunc_sig = (typeof(f), typeof(x), typeof(y), Constant{typeof(z)}, Cache{typeof(c)})

# Approach 1
@rule_from_DifferentiateWith(target_backend, ::Tfunc_sig, source_backend)

# Approach 2: equivalent but in two steps
dw = DifferentiateWith(::Tfunc_sig, source_backend)
@rule_from_DifferentiateWith(target_backend, dw)