Taking Complex Autodiff Seriously in ChainRules

Specific Domains
, ,
62

Okay, so after some very productive and stimulating conversation, this recent push has resulting in ChainRulesCore.jl/pull/167 (thanks @ettersi!) and ChainRules.jl/pull/196 which I must say feels like we’ve come up with a really great set of shared definitions and conventions.

The result of @ettersi’s now merged PR can be viewed in all of it’s LaTeX glory here:
https://www.juliadiff.org/ChainRulesCore.jl/dev/complex.html

Perhaps next up, we can lift some ideas from here: Complex Differentiation · Zygote such as the example showing how to get the Wirtinger derivatives, \partial \over \partial z and \partial \over \partial \bar{z} from the Jacobian, i.e.

using ChainRules
function wirtinger_rrule(f, z)
    _, pullback = rrule(f, z)
    du, dv = pullback(1)[2], pullback(im)[2]
    (du' + im*dv')/2, (du + im*dv)/2
end

wirtinger_rrule(abs2, 1 + im)

#+RESULTS:
: (1.0 - 1.0im, 1.0 + 1.0im)

function wirtinger_frule(f, z)
    du_dx, dv_dx = reim(frule((Zero(), 1),f,z)[2])
    du_dy, dv_dy = reim(frule((Zero(),im),f,z)[2])

    (du_dx + dv_dy + im*(-du_dy + dv_dx))/2, (du_dx - dv_dy + im*(du_dy + dv_dx))/2
end

wirtinger_frule(sin, 1 + im)

#+RESULTS:
: (0.8337300251311491 - 0.9888977057628651im, 0.0 + 0.0im)
8 Likes