Description
I increasingly feel like all the heavy lifting should be done by backends, and AbstractDifferentiation.jl should just define common names.
Whenever we do something that a backend does not support, or when we recode it instead of letting the package do it, we introduce correctness, readability or performance issues. Some examples:
- additional closures with type instabilities (Fix type instabilities #109, Regression in type inferrability #121)
- hard-to-parse macros (Get rid of ugly macro? #123)
- vicious tuples (
pushforward_function
andpullback_function
are confused by tuples vs single input #99, Use multiple arguments instead of a tuple for pushforward and pullback function? #53) - no exploitation of optimized backend code (
AD.jacobian
much slower thanZygote.jacobian
#54)
Moreover, the complexity of the package also makes it very hard to push Enzyme support over the finish line (#84), which is arguably the biggest missing feature.
@adrhill and I are working on a prototype to play with this idea:
https://github.com/gdalle/DifferentiationInterface.jl
It may or may not be merged into AbstractDifferentiation, depending on what the team here thinks. The main caveat I see is that higher order cannot be well supported with an empty shell, but for gradients and JVPs/VJPs it should be close to optimal.