Custom rules with broadcasting

[Antomek]

Custom rules with broadcasting

[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.