Complex numbers

[willtebbutt]

This was opened as a response to the discovery in this Zygote PR that our handling of complex functions isn't currently Zygote-compatible when they're not analytic.

The issue appears to stem from the need to conjugate the input / output to rrules when used in Zygote. This is a direct consequence of chain rules currently assuming that all functions for which rules are implemented are holomorphic, as in this case simply conjugating the input and output to the rule makes total sense and does the correct thing.

What needs to happen is:

• ChainRulesCore to get a @holomorphic or @analytic macro for handling scalars involving complex numbers that does the conjugating for you
• ChainRules to use this to implement the existing holomorphic functions
• Careful testing of scalar functions with complex inputs in ChainRules (I'm assuming the lack of this is how stuff got through in the first place undetected)
• A short PR to Zygote that simply stops conjugating things

[nickrobinson251]

xref #23 #40 JuliaDiff/Capstan.jl#1 FluxML/Zygote.jl#142

[oxinabox]

I am not sure the conjugation and holomorphicness are related.
AFAICT Zygote just always stores and works with the conjugate of the complex gradient.
Just at all times, including returning it.
We work with the actual gradient.
At least it think the gradient we have is the actual one, since we agree with finite differencing where as Zygote gets the conjugte of the result that FiniteDifferences.jl gives

(cc @MikeInnes)

[willtebbutt]

Hmm okay, then I'm slightly confused. The following:

julia> x = 5.0 + 4.0im
5.0 + 4.0im

julia> y = 4.0 + 3.0im
4.0 + 3.0im

julia> z, back = Zygote.pullback(+, x, y)
(9.0 + 7.0im, Zygote.var"#36#37"{Zygote.var"#1587#back#606"{Zygote.var"#604#605"{Complex{Float64},Complex{Float64}}}}(Zygote.var"#1587#back#606"{Zygote.var"#604#605"{Complex{Float64},Complex{Float64}}}(Zygote.var"#604#605"{Complex{Float64},Complex{Float64}}(5.0 + 4.0im, 4.0 + 3.0im))))

julia> back(1.0 + 1.0im)
(1.0 + 1.0im, 1.0 + 1.0im)

suggests that Zygote at least returns the cotangent, as opposed to the conjugate of the cotangent. Am I missing something @oxinabox ?

edit: it might be that I should write something up about the basics of AD + complex numbers if no one has already written anything.

[oxinabox]

I am not sure about anything to do with complex number AD. so a document would be good

[sethaxen]

I'd love to see that. Also how it connects to AD through functions with real inputs and outputs but complex intermediates.

[simeonschaub]

The analytical complex derivative is defined as ∂/∂z = ∂/∂x - im*∂/∂y. What Zygote does is it returns ∂/∂x + im*∂/∂y, which is equivalent to the conjugate of the analytical derivative if the output is eventually real, which is always the case in Zygote.
@willtebbutt I believe reverse mode always pulls back one forms, which live on the cotangent space, but we might not always just want to use Adjoint{<:Vector} to represent them, as that doesn't generalize well to functions operating on higher dimensional arrays, structured derivatives (e.g. Composite) and also complex numbers, if you want to identify them with the linear functionals on R^2N, as the complex inner product on C^N is completely different to the real inner product on R^2N.
If you are interested, we could also discuss this in a call.

[willtebbutt]

The analytical complex derivative is defined as ∂/∂z = ∂/∂x - im*∂/∂y.

🤦 you're right of course. I clearly need to revise some of the basics lol

If you are interested, we could also discuss this in a call.

Definitely interested. I've dropped you a message on slack.

[sethaxen]

At least it think the gradient we have is the actual one, since we agree with finite differencing where as Zygote gets the conjugte of the result that FiniteDifferences.jl gives

Does ChainRules agree with FiniteDifferences? e.g. is this the expected behavior (on FDv0.10.0)?

julia> using FiniteDifferences

julia> j′vp(central_fdm(5, 1), sin, 1.0 + 0.0im, 2.0 + 3.0im)
(-4.18962569096891 + 9.10922789375644im,)

julia> j′vp(central_fdm(5, 1), sqrt, 1.0 + 0.0im, 2.0 + 3.0im)
(0.23216272632545484 + 0.1242497204556949im,)

julia> using ChainRules

julia> ChainRules.rrule(sin, 2.0+3.0im)[2](1.0+0.0im)
(Zero(), -4.189625690968807 - 9.109227893755337im)

julia> ChainRules.rrule(sqrt, 2.0+3.0im)[2](1.0+0.0im)
(Zero(), 0.2321627263254075 - 0.12424972045565168im)

[oxinabox]

No that is not the expected behavour. At least not the behavour, I expected.
oh dear, clearly I do not know what i am talking about.
I thought I always had to conjugate Zygote's output to agree with FIniteDIfferences...
and I always had to conjugate ChainRules to agree with Zygote

[sethaxen]

I thought I always had to conjugate Zygote's output to agree with FIniteDIfferences...

For the complex rules I've checked, it seems like Zygote's adjoints generally agree with FD's j′vp's. e.g.

julia> A, B, C̄ = randn(ComplexF64, 3, 3), randn(ComplexF64, 3, 3), randn(ComplexF64, 3, 3);

julia> all(Zygote.pullback(*, A, B)[2](C̄) .≈ j′vp(central_fdm(5, 1), *, C̄, A, B))
true

julia> all(Zygote.pullback(inv, A)[2](C̄) .≈ j′vp(central_fdm(5, 1), inv, C̄, A))
true

[oxinabox]

@willtebbutt can you run this whole thing down with people and come back with an explination of what is it is exactly that we do differently to Zygote and if we are just wrong?
I won't have time to look at this this week

[MasonProtter]

I'd say this is expected. The \prime in j′vp (as well as the word 'adjoint' in the docstring) suggests that j′vp(f, v, x) is adjoint((D(f))(x))*v where I take D(f)(x) to mean the derivative of f evaluated at x.