Custom rules with broadcasting

[Antomek]

Hello all,

Thank you for all your work on ChainRules, it is shaping up to be very useful.
I recently wrote my first frule/rrule for some (scalar) function, without using the @scalar_rule macro.
It worked fine for 1D case, but I noticed it did not work with broadcasting (as in, my custom rules were never used).

Is there a straighforward way of getting custom rules to play nicely with broadcasting, or did I miss something obvious?
It might be worth adding something to the docs regarding this, since broadcasting seems to be central Julia concept.

Many thanks!

[mcabbott]

Assuming you mean Zygote, its broadcasting at present prefers to use ForwardDiff internally, if possible. Thus it will not see rules defined using either Zygote itself, or using anything from ChainRules.

What you can do it write a rule for broadcasted(f, x), e.g. https://github.com/FluxML/NNlib.jl/blob/master/src/activations.jl#L880 or https://github.com/JuliaDiff/ChainRules.jl/pull/586/files#diff-69fb16cdd94960c65024f9dfe644deae5023a2af02d68f26c0855cf538d2ca0aR53 . If, like these, your rule returns broadcast(f, x) then it will un-fuse the broadcasting. Which at present Zygote also does.

This is not well documented, unfortunately. Partly because it's not so clear what the best strategy is. @scalar_rule now defines also derivatives_using_output which should be useful for handling broadcasting.

[oxinabox]

This is basically a zygote issue which is not actionable further here, so I am closing this for now.