Higher-Order functions: Low-hanging fruit

[willtebbutt]

We've yet to properly figure out how to make ChainRules play nicely with higher-order functions, see JuliaDiff/ChainRulesCore.jl#68.

However, this doesn't stop us from implement a load of rules for functions that ChainRules knows how to handle explicitly. For example, we have a rule for sin, so we can implement map(sin, x) ourselves for carefully chosen types of x.

A pre-requisite for this is JuliaDiff/ChainRulesCore.jl#127 which would let us get hold of all of the rules that we've implemented. It'll be necessary to be able to filter them for appropriateness so that we know what they're applicable to, in particular we need to know

• the arity of the function to which the rule applies
• the element types to which it applies

This would be a big win for e.g. Zygote where this is done, but with a piece-meal approach.

[oxinabox]

do we need to know the return type also?

[willtebbutt]

I don't think so 🤷 . Is this a particular concern?

[oxinabox]

we don't know return types.
But if we don't need to, it is easy

[mzgubic]

We can probably close this now since JuliaDiff/ChainRulesCore.jl#363 is merged?