Some thoughts on Wirtinger

[oxinabox]

This comes up in discussion on Zygote
FluxML/Zygote.jl#142
FluxML/Zygote.jl#291

Mike/Zygote doesn't like Wirtinger.
Which I personally can get behind, since it is well outside the math I use.
And anything to do with them is a maintaince issue because
I think of the people the contribute to this repo only @simeonschaub and @jrevels properly get them,
(both in terms of the how, and the why)

Zygote never wants to see a Wirtinger differential result.
I think though that might mean also that Zygote wants to treat nonholomorphic functions as nondifferentiable?

Would it make sense to define

function extern(w::Wirtinger)
    iszero(extern(wirtinger_conjugate(w))) || throw(DomainError("Not holomorphic"))
    return extern(wirtinger_primal(w))
end 

Or the equiv inside Zygote.

Could we have a seperate rrule/frule
for Wirtinger that is allowed to return Wirtinger
and normal use of rrule and frule don't.
Then general_rrule and general_frule would be fallback to rrule and frule if there is no special overload for the Wirtinger case defined?

[simeonschaub]

With #29 merged, WirtingerRule should only be returned if the function is actually non-holomorphic and there's no way to represent the derivative without a WirtingerRule. There might still be single points in a function's domain where wirtinger_conjugate is zero (eg. abs2(0+0im)), but does it really make sense to say, that a complex derivative exists at single, isolated points? So for Zygote, the current behavior, that an error is always thrown when extern is called on a Wirtinger, should be adequate. In general, it might make sense to only return Wirtinger if explicitly requested, and return nothing otherwise, like you propose, since this would be closer to the mathematical notion, that these functions are not complex differentiable. On the other hand, it should be easy enough to check if a WirtingerRule is returned, and Wirtinger will throw an error, as soon as you call extern on it, anyways.

[jrevels]

I don't have time these days to be a super active maintainer (@simeonschaub has been doing great work getting the prototypical Wirtinger implementation into a usable state!), but we should just keep in mind the actual feature we want to provide here: enable downstream tools to get useful gradients when optimizing models containing non-holomorphic components (cc @ssfrr, who originally had some interesting applications here).

The goal was to provide this feature in a manner that "just works" without real-valued users ever needing to involve themselves with the feature + allow complex AD to be first-class w.r.t. supporting tensor primitives + compose seamlessly with real-valued primitives, etc. The Wirtinger approach is what popped out in the pursuit of that goal. If there are specific problems (which I'm sure there are, this approach hasn't really been "battle-tested" yet), would be great to highlight them and brainstorm solutions and/or replacement implementations for this feature

[oxinabox]

enable downstream tools to get useful gradients when optimizing models containing non-holomorphic components (cc @ssfrr, who originally had some interesting applications here).

I guess I would like to understand that much more.

Unlike all other (non-error related) Differentials
an a Wirtigner can not be (logically) subtracted from something of the same type as spawned it.
So I don't see how it can be used in optimizeion,
can you recommend some reading? Maybe a blog or wikipedia page?

[simeonschaub]

The Optim.jl docs have a section on complex optimization here. The g they define there is nothing other than our wirtinger_conjugate. See also #23, where I proposed how to better support exactly these use cases. The full Wirtinger is still needed, if there are intermediate non-holomorphic parts, though.

[ssfrr]

@oxinabox There are some good resources and links to other relevant discussions collected here

Let's take some protypical function f: C^N->R that's composed of some other primitive functions that have rules defined.

If all you care about is gradients then you can decompose f function inputs into the real and imaginary pairs, treat it as a function R^(2N)->R and get a meaningful gradient.

The problem is that this function doesn't have a derivative(fn1), so if you ask the AD framework for its derivative you shouldn't just get a number that you can blindly (wrongly) move forward with.

@jrevels and @simeonschaub can correct me if I'm wrong here (I have a reasonable handle on the math but not on the AD implementation), but I think the main motivations for Wirtinger are:

1. Keep track of whether composite functions are holomorphic/antiholomorphic in the type system, based on how the primitive functions compose
2. Allow diff rules to be written that don't require splitting up the input into the real and imaginary parts.
3. Knowing the structure of the primitive functions and their compositions allows some optimizations

Point (1) is actually a little less compelling to me now than it used to be, because I think you still end up wanting to do a numerical check for Cauchy-Riemann to see if the composed function is holomorphic. For example, f(z) = real(z) + im*imag(z) is non-holomorphic according to the type system, but if you check Cauchy-Riemann numerically it is holomorphic (fn2).

IIRC point (2) is one of the things that got @jrevels excited about this in the beginning. If you're writing a rule where the input to the function is some large Vector{ComplexF64} it would be nice to keep Complex as an atomic unit, rather than splitting it into two Vector{Float64}s.

(3) is also pretty cool, e.g. for h(z) = g(f(z)), if you know that wirtinger_primal(df) = Zero() and wirtinger_primal(dg) = Zero(), then you know that wirtinger_conjugate(dh) = Zero() (fn3). As with all optimizations however, I'd want to verify in the wild that they actually are useful.

The general feeling I had (and I think @jrevels though I don't want to speak for him) is that the Wirtinger stuff will be particularly beneficial as we start writing diff rules for array-valued primitive functions. At the time he didn't get far enough into the implementation to prove that out - @simeonschaub maybe you have some insight here?

Footnotes

(fn1): Where I'm defining the derivative as a local linear approximation Df of the function f such that f(z+ϵ()) ≈ f(z) + Df*ϵ(), and ϵ() could be any small complex perturbation: ϵ() = (rand(ComplexF64) - 0.5 - im*0.5)*0.001

(fn2): One kinda nice thing about Wirtingers is that the holomorphism check is just wirtinger_conjugate(df) ≈ 0. If you split into real and imaginary parts it is real(df))/real(dz) ≈ imag(df)/imag(dz) && imag(df))/real(dz) ≈ -real(df)/imag(dz) but computers don't really care.

(fn3): This is equivalent to saying that if f and g are antiholomorphic, then g ∘ f is holomorphic

[simeonschaub]

Regarding point (1): I don't know how representative the f(z) you pick here actually is. Having good native support for complex differentiation should actually mean, that you shouldn't really have to write code like this. If you separate real and imaginary parts like this and in the end recombine them and find the Wirtinger conjugate is zero most of the time, that should be an indication that you can rewrite your code using only z as a whole. Still, doing this check if the result is a Wirtinger also shouldn't really hurt.
Regarding point (2) and arrays in general, that's something that still needs quite a bit of work. One problem is, that in many cases you will get a differential of type Array{Wirtinger{Complex,Complex}}, but in other cases, if some of your input is an array, returning Wirtinger{Array{Complex},Array{Complex}} might make more sense. The latter one could also be converted to the first kind in the constructor of Wirtinger, but I think this decision has to be made, once this is actually used in the real world by package authors.
Point (3) actually already works very nicely, just by the way Zero propagates. Although we should probably have some kind of fall-through for Wirtinger(df_dz, Zero()) to just df_dz, that doesn't trigger #9 again.

[oxinabox]

@ssfrr thanks, I will return to this later and read all these in detail next time I need to understand why we do this

[oxinabox]

http://www.ekinakyurek.me/complex-derivatives-wirtinger/

[oxinabox]

I am back on to the "I think we are wrong to do Wirtinger" line.
Had a long talk with @awf re-this and many things.

The major problem with Wirtinger is it violates the principal behind
#8
That the type of a differential returned from a pushforward should be addable to the function output.
And from pullback it should be addable to the input.

The input in this case is a complex number.
Which is a pair of things. 2 things.
Wirtinger is a pair of complex numbers,
so it is 4 things.

Why is it 4 things?
I think because it is a representation of the Jacobian. No?

Where as pullback is normally returning a gradient. (This is the whole vjp thing?)

And we mostly haven't run into big problem because we only have scalarrules for Wirtinger.

(Also the time we tried to do a thing with it we did run into problems. Multiplication is hard, externing is hard.)

My proposal is that after #30 we move it to its own package that extends ChainRules.
We can keep that working as a bit more experimental, while we focus on finishing up and maturiting ChainRules.

And ChainRules returns the structured derivative (#8) (or equiv a Complex number) if defined or we say it doesn't exist.

Once things are more developed we can return and reintegrate WirtingerChainRules.jl if at that stage it seems the best plan.

[awf]

And in particular, I'm thinking of the situation where you want to minimize some quantity of interest.
Just as in the discussions linked above, but perhaps more colourfully: imagine you're designing a widget with lasers in it, and you want to optimize the goodness of the widget. One thing you might want is that gradient descent "just works"...

struct LaserThingy
    phasor::Complex
    speed::Real
end

# Measure goodness of LaserThingy
goodness(laser::LaserThingy) = abs(log(1im + laser.phasor)) + laser.speed^2

# Make an initial estimate
l = LaserThingy(1-2im,3.1)      

# And optimize it using gradient descent...
using Zygote

import Base: *,+
# Generically, we should be able to add Ts to DNamedTuples derived from Ts...
+(a::LaserThingy, b::NamedTuple{(:phasor, :speed),Tuple{Complex{Float64},Float64}}) = LaserThingy(a.phasor+b.phasor,a.speed+b.speed)  
# … and to scale DNamedTuples
*(s::Float64, x::NamedTuple{(:phasor, :speed),Tuple{Complex{Float64},Float64}}) = (phasor=s*x.phasor,speed=s*x.speed)     

eta = 0.001 
for k in 1:100
    global l
    g = Zygote.gradient(goodness, l)[1]
    l = l + -eta*g
    println(goodness(l), "   ", l)
end 

[ssfrr]

I agree that users of Zygote should never see a Wirtinger. I would imagine that Zygote.gradient would perform AD on the users code, and then if it gets a Wirtinger it just returns wirtinger_conjugate(g), which turns out to be a gradient vector that you can add to the functions inputs and gradient ascend to your hearts content. This is valid whether or not the user's composite function is holomorphic. The whole Wirtinger thing seems like it's useful internal to the AD implementation because they compose nicely and retain information about holomorphism.

One of the ideas that drove this design in the first place though is that gradients aren't the only thing people want to use AD for - people also want to compute derivatives. If the user calls Zygote.derivative then I'd propose that Zygote performs AD as before, but then uses Cauchy-Riemann to assert that the provided function is actually holomorphic (more details on this in my message above). If the function is holomorphic (wirtinger_conjugate(g) == 0) then Zygote can return wirtinger_primal(g)).

One core point here is that derivatives != gradients. There are cases where one exists but not the other, and in the cases that they both exist they are adjoint (complex conjugate) to each other.

[simeonschaub]

I want to add that Zygote's definition of a complex gradient generally only works for reverse-mode AD, because for forward-mode, you will need to represent intermediate derivatives of non-holomorphic C->C functions, even if you just want to differentiate C->R functions. Since the main purpose is to have both native support for complex functions as well as support for interoperable mixed-mode AD, I don't really see a way we could do this without a Wirtinger type.

[oxinabox]

How do people feel about my proposal to move Wirtinger out to
WirtingerChainRules.jl
where we can experiment with the idea more,
and develop it properly,
e.g. get some definitions in that return Arrays of Wirtingers etc.
Then merge it back in later.

Right now, I would like to be pushing towards ChainRules being mature enough to use in practice,
and the Wirtinger stuff is causing me problems, perhaps because i don't understand it.
And it doesn't work well right now.
Recall the problems last week:

Lyndon White:ox: Sep 12th at 5:33 PM

@simeon Schaub did I screw something up:

using ChainRules gives result of
#Δx = Wirtinger{Complex{Int64},Complex{Int64}}(0 - 4im, 4 + 0im)
But Zygote and FiniteDifferences think it is 4
complex_function_example.jl 
using ChainRules # with ChainRulesCore #30 and matching ChainRules PR
using FiniteDifferences
using Zygote
​
function foo(x)
    a = x * (1+1im)
    b = abs2(a)
    c = real(b)
    return c
end
​
foo(1)
​
_, pb = Zygote._forward(foo, 1)
_, x_zy = pb(1)
@show x_zy
#x_zy = 4 + 0im
​
x_fd = central_fdm(5,1)(foo, 1) 
@show x_fd
#x_fd = 3.9999999999999765
​
######
​
x = 1
a, a_pb = rrule(*, x, (1+1im))
b, b_pb = rrule(abs2, a)
c, c_pb = rrule(real, b)
​
Δc = 1
_, Δb = c_pb(Δc)
_, Δa = b_pb(Δb)
_, Δx, _ = a_pb(Δa) 
​
@show Δx
#Δx = Wirtinger{Complex{Int64},Complex{Int64}}(0 - 4im, 4 + 0im)

Lyndon White:ox: 8 days ago

Aside: writing out chainrules invokations like this is pretty good as a thing to do in the docs

Lyndon White:ox: 8 days ago

Oh, is this
JuliaDiff/ChainRules.jl#95

Lyndon White:ox: 8 days ago

even if I add the robutified version of #95

function ChainRulesCore.rrule(::typeof(*), x::Number, y::Number)
    Ω = x * y,
    function times_pullback(ΔΩ)
        ∂x = differential(typeof(x), ΔΩ * y)
        ∂y = differential(typeof(y), x * ΔΩ)
        return (NO_FIELDS, ∂x, ∂y)
    end
end

I still get a result of Δx = 2+2im
Also sent to the channel

Lyndon White:ox: 8 days ago

this is really annoying
The real can be dropped since it acts as an identity.
Swap the abs2 for an abs and we still disagree with zygote and FDM
Also sent to the channel

Simeon Schaub 8 days ago

This is something that stunned me at first too, and I had a very enlightening discussion with Mike Innes about this. The short answer is that zygote uses a completely different definition for complex gradients, which works very well for non-holomorphic C->R functions, so neither one of them is actually wrong. I'll provide a proper explanation tomorrow.

Simeon Schaub 7 days ago

Zygote actually defines the pullback of a ℂ->ℂ function f(x + imy) as ∂(Re(f))/∂x + im∂(Re(f))/∂y. This is equivalent to 2*∂(Re(f))/∂z̅ = ∂(f + f̄)/∂z̅ = ∂f/∂z̅ + conj(∂f/∂z). From this, it is easy to see, that for holomorphic functions, the pullback is just the conjugate of the primal derivative and for ℂ->ℝ functions, this is just twice the conjugate derivative of f. One can then calculate the pullback of f∘g and one finds that iff the outer function f is equal to its conjugate, so its output is always real, the chain rule applies, which means to get the pullback of f∘g, one can just multiply the pullbacks of f and g. One major implication of this, is that this only works for reverse-mode AD, since for forward-mode, the outer functions' output does not always have to be real. This also makes it quite expensive, if you actually want to get the full Wirtinger derivative of a ℂ->ℂ function, since you have to calculate the whole backwards pass twice. To get the Zygote behavior using ChainRules, it would be very useful to have the mechanism I describe in ChainRulesCore.jl#23. I'll try to make a PR to the Zygote docs soon trying to explain this relationship a bit better.

Lyndon White:ox: 7 days ago

Can one define extern_zygote(::Wirtinger)
to do that operation?
(or error if impossible) (edited)

Simeon Schaub 7 days ago

extern_zygote(x::Wirtinger) = wirtinger_conjugate(x) + conj(wirtinger_primal(x)) should work, but we could do this more efficiently with #23

Lyndon White:ox: 7 days ago

hmm I am still getting a complex number

Lyndon White:ox: 7 days ago

0.7… + 0.7 im
Where as Zygote and Finite differences give 1.4…

Simeon Schaub 7 days ago

For what function?

Lyndon White:ox: 7 days ago

Well for orginal example at top,
Zygot/Finite difference gives 4
where as
chainrules with that extern_zygote gives 2+2im
x = 1
a, a_pb = rrule(*, x, (1+1im))
b, b_pb = rrule(abs2, a)

Δb = 1
_, Δa = b_pb(Δb)
_, Δx, _ = a_pb(Δa)

@show Δx
wirtinger_conjugate(Δx) + conj(wirtinger_primal((Δx)))

#Δx = Wirtinger{Complex{Int64},Complex{Int64}}(0 - 4im, 4 + 0im)
@show differential(typeof(x), Δx)
(edited)

Simeon Schaub 7 days ago

Ah, that's probably because our derivative for complex * is wrong

Simeon Schaub 7 days ago

See ChainRules.jl#95

Lyndon White:ox: 7 days ago

That is with

function ChainRulesCore.rrule(::typeof(*), x::Number, y::Number)
Ω = x * y,
function times_pullback(ΔΩ)
∂x = ΔΩ * y
∂y = x * ΔΩ
return (NO_FIELDS, ∂x, ∂y)
end
end

Lyndon White:ox: 7 days ago

Actually if I remove that,
and go back to using ' in rrule(*
and I z_extern between every chain’ing.
(which isn’t too unreasonable since often need to extern between calling pullbacks)
Then we do get same answer as Zygote/Finite differencing (edited)

Simeon Schaub 7 days ago

Good catch! We need a different rule for * with Wirtinger

Simeon Schaub 7 days ago

Try
function times_pullback(ΔΩ::Wirtinger)
∂x = Wirtinger(wirtinger_primal(ΔΩ) * y, wirtinger_conjugate(ΔΩ) * conj(y))
∂y = Wirtinger(x * wirtinger_primal(ΔΩ), conj(x) * wirtinger_conjugate(ΔΩ))
return (NO_FIELDS, ∂x, ∂y)
end
(edited)

Simeon Schaub 7 days ago

We should just use @scalar_rule for this, as this does this kind of Wirtinger chaining already. I don't think we really want to support non-commutative Numbers for now.

Lyndon White:ox: 7 days ago

Yessss, this works
🎉
1

Lyndon White:ox: 7 days ago

But if I can’t z_extern in between pullbacks, or i get 4im
Which I would like to do,
since in general externing between pullbacks is required, since not all functions are implemented on AbstractDifferentials.
(and I kinda think that is OK)

Simeon Schaub 7 days ago

I wonder whether we should disallow complex scalar multiplication of Wirtinger altogether, because the chaining is different for reverse mode vs forward mode

Simeon Schaub 7 days ago

But if I can’t z_extern in between pullbacks, or i get 4im
Which I would like to do,
since in general externing between pullbacks is required, since not all functions are implemented on AbstractDifferentials.
(and I kinda think that is OK)
The problem is that wirtinger_[primal|conjugate] aren't recursive. I'll make a PR to fix that.

Simeon Schaub 7 days ago

Ah, sorry. I misunderstood your problem. I think you need to pass Wirtinger(conj(Δa_externed)/2, Δa_externed/2) to a_pb but I could be wrong.

Simeon Schaub 5:19 PM

I think we need a function chain in ChainRulesCore taking two differentials, which usually just falls back to multiplication, but if any of the arguments is a Wirtinger, treats the first argument as the partial derivative of the outer function and the second as the derivative of the inner function. We could then disallow multiplication of Wirtinger with complex differentials, since for these cases it is important which differential is that of the inner/outer function. This is why * produced the wrong derivative if used together with Wirtinger. This could potentially also simplify the logic in @scalar_rule, although I haven't looked into that too thouroughly yet.

So for now, I propose it goes into its own package, and we let it develop at its own pace.
Mean while ChainRules.jl and ChainRulesCore.jl can make sure everything we do, we do well.
and if it comes to it we accept that there are somethings we simply can't do.
(until we remerge the WirtingerChainRules)

This to me feels better to let things go at their pace.

[simeonschaub]

If you don't want to bother with `Wirtinger`, can't you just error as soon as one is encountered and we'll mark complex differentiation as highly experimental? It seems like a maintenance nightmare to me, to have two different versions (and another for `ChainRules`, too), and having to keep both up to date. We would have to move all of the `Wirtinger` stuff out of both repos and merge them back in again later, when there are probably a bunch more changes. `Wirtinger` is also deeply integrated into `@scalar_rule`, so it won't be trivial to have just a package on top of `ChainRulesCore`. I don't really feel like having `Wirtinger` here does much harm, as it will error anyways, any time you try to `extern` it. We should, for now, document properly though, how broken it is.

[oxinabox]

I don't think it would be that hard to split it out.

And I think while Wirtinger is experimental,
we could have all the Wirtinger stuff,
both rules and Core in one package,
which would make it faster to experiment and improve that.

Wirtingeris also deeply integrated into@scalar_rule`

Not so much as it might seem.

The current definition of scalar_rule would move to be @wscalar_rule in the Wirtinger package,
and the one that left behind would have the propagate_wiringer removed.

---

Using Wirtinger in AD is a fairly novel idea, and maturing it is going to take much longer than anything else in the package.