After some reflection, I think the current Zygote behaviour is the right way to go. It supports pretending everything is holomorphic, but also defines the rules correctly such that you can get out the full Jacobian if you want / need. Following @sethaxen’s link to the Zygote docs, we see that the way it works is that if you do
using Zygote
y, back = Zygote.pullback(abs2, 1.0 + im)
julia> back(1), back(im)
((2.0 + 2.0im,), (0.0 + 0.0im,))
this is the full Jacobian. back(1) is asking for (J*[1, 0])' and back(im) is asking for (J*[0, 1])', and we can then form the Wirtinger derivatives via
du, dv = back(1)[1], back(im)[1]
(du' + im*dv')/2, (du + im*dv)/2
To my surprise, this actually works on @oxinabox’s Chainrules branch of Zygote, even though ChainRules specifically defines
@scalar_rule abs2(x) 2x
Is there something I’m misunderstanding here or is Zygote somehow skipping that chain-rule on this branch?