Zygote only sees untyped, pre-inference IR (think @code_lowered
or @code_warntype
without type information). It simply does not have enough information to do anything beyond the simplest, surface level transformations.
That said, it looks like we can force Zygote to evaluate pullbacks for mutating functions even if the result isn’t used. See rrule for fill! by CarloLucibello · Pull Request #521 · JuliaDiff/ChainRules.jl · GitHub.
One issue with doing this via a wrapper is that there are hundreds of existing ChainRules rrules (spread over numerous packages) out there which take array arguments. Making all of those aware of the wrapper type and able to check for validity/poison would be a massive effort, hence the exploration of less distruptive alternatives like writing NaNs. It’s quite possible this would require major changes to ChainRulesCore’s interface, as discussed in Ability to specify different rules based on what combinations of inputs are actually being used · Issue #452 · JuliaDiff/ChainRulesCore.jl · GitHub, mutating calls · Issue #242 · JuliaDiff/ChainRulesCore.jl · GitHub, etc.