Ah, yes I missed that, it is an interesting way to do it. I guess we could do the gradient the way they do it but also supply a jacobian that gives the real deal.
I’m not sure I agree with
though, but I suppose if there’s a place to enforce holomorphic assumptions, gradient would be an okay place to do it.
Distinguishing between gradients and derivatives is a very sensible thing to do.
For illustration, consider |x|:
Both (co-)gradients as steepest-ascent directions and derivatives for Taylor-like expansions stem from the same thing in CR calculus, though notations vary. You can write the multidimensional “Taylor” expansion f(x+\delta) = f(x) + \bar{\delta}\cdot\nabla_\bar{x} f + \delta\cdot\nabla_x f + O(\delta^2), and if f is real-valued then \nabla_x f= \overline{\nabla_\bar{x} f} and \nabla_\bar{x} f = (\frac{\partial f}{\partial \bar{x}_1}, \ldots, \frac{\partial f}{\partial \bar{x}_n}) is the steepest-ascent direction in x. (If f is complex-valued then you can’t talk about “ascent” or “descent.”) In that real-f case, however, since f(x+\delta) = f(x) + 2\Re[\bar{\delta}\cdot\nabla_\bar{x} f] + O(\delta^2), you need to pass 2\nabla_\bar{x} f as the “gradient” to any code expecting the gradient in terms of the real and imaginary parts (e.g. an optimization routine that only knows about real vectors).
Since CR calculus is widely used and fairly standardized, I don’t think we need to get into metaphysical arguments about whether gradients are “sensible” for non-holomorphic functions. They are undeniably useful, and the only question in my mind is what interface to present.
Since CR calculus is widely used and fairly standardized, I don’t think we need to get into metaphysical arguments about whether gradients are “sensible” for non-holomorphic functions. They are undeniably useful, and the only question in my mind is what interface to present.
I am starting to realise that non-holomorphic functions are important in some fields (mostly electrical engineering, afaict), but so are holomorphic functions in other fields (e.g. complex analysis and approximation theory, which is where I come from). I have a feeling that the discussion so far heavily favours the former. I can imagine that this is somewhat sensible since there are many more electrical engineers than there are people working on complex analysis, but I do hope we can find a framework which supports non-holomorphic derivatives without breaking (or at least severely complicating) the holomorphic case.
The challenge then seems to be that for f : \mathbb{C} \to \mathbb{C}, we want f' : \mathbb{C} \to \mathbb{C} in the holomorphic case, but f' : \mathbb{C} \to \mathbb{C}^2 or f' :\mathbb{C} \to \mathbb{R}^{2 \times 2} or any other variation of this theme in the non-holomorphic case. I still believe that the “default” (e.g. in an expression like f'(x)) should be to assume that f(x) is holomorphic and return the holomorphic derivative. If not, I cannot write generic code for dealing with a function f(z::T) -> T where T <: Number, since then f'(z) would be of type T if T <: Real but f'(z) would be something else if T <: Complex.
Actually, Autograd’s convention, Wirtinger derivatives and holomorphic derivatives are equivalent in all cases where the various derivatives are defined. So the only thing that is needed is a system for turning on and off whether to compute and return \tfrac{\partial f}{\partial \bar z}.
As far as I can tell, the only clean way of implementing this is to have frule and rrule return \nabla_z f, and then provide separate wirtinger_frule and wirtinger_rrule functions which return (\nabla_z f, \nabla_{\bar z} f).
Apologies, I should have been more precise – I definitely wasn’t trying to engage in a metaphysical debate 
. I was just recapitulating the point that you generally need four real numbers to talk about the derivative of a \mathbb{C} \to \mathbb{C} function f at a point, rather than two.
I am definitely not debating is the need for non-holomorphic Chain Rules – the need for such rules is very clear. Moreover I completely agree with the idea that an AD tool really has to assume that everything is non-holomorphic by default, while potentially allowing a user to specify that they believe the function they’re asking for the pushforward / pullback of if holomorphic, and that they would therefore like an answer than assumes that (such an assumption probably won’t be able to avoid much computation. As has been discussed previously in this thread, it’s perfectly possible for a holomorphic function to be a composition of non-holomorphic functions).
@ChrisRackauckas made the central point here, which I’m going to rephrase slightly: it’s helpful to forget about derivatives and gradients for now, and focus on pushforwards / jacobian-vector products and pullbacks / vector-jacobian products. These are the fundamental operations with which an AD tool / set of rules for an AD tool (ChainRules) is concerned, as gradients / derivatives follow from these. The criteria for choosing how they should be defined should probably consider the implications for special cases of interest, such as for holomorphic \mathbb{C} \to \mathbb{C} functions, and the interface that an AD tool implement for the gradient of a \mathbb{C}^N \to \mathbb{C} or \mathbb{C}^N \to \mathbb{R} (implemented in terms of a pullback / vector-jacobian product).
As has already been discussed, one way of tackling these is just by treating complex numbers and their tangents as length-2 vectors, and just naively deriving everything from there. Autograd / JAX don’t appear to do this completely naively, which is why I linked to their conventions. So my question here is whether or not these are sufficient / ideal, or whether there is room for improvement. Is there some tangible advantage to the choices that JAX / Autograd make? From what Chris discusses above, perhaps it doesn’t really matter 
.
Again, I’ve not had time to properly think through everything here, but it’s important that the discussion is centred around how the operations that ChainRules / AD tools actually implement should be defined.
In a case \mathbb{C}^m \to \mathbb{R}, intermediate functions will generically be of the form \mathbb{C}^p \to \mathbb{C}^s. The pullbacks will take \mathbb{C}^s \to \mathbb{C}^p, but these don’t need any special structure; they should do the same thing as the equivalent \mathbb{R}^{2s} \to \mathbb{R}^{2p}. In such a case, one shouldn’t need a complex Jacobian, and many reverse-mode rules can be implemented the same way for \mathbb{C}^p \to \mathbb{C}^s as for \mathbb{R}^p \to \mathbb{R}^s (see e.g. Zygote’s rules).
Perhaps it’s my bias as someone who only needs Jvps for \mathbb{R}^m \to \mathbb{C}^s \to \mathbb{R}, but it would be a shame if the result of this was that rules that are perfectly fine for such complex differentiation need to be completely rewritten to support this more general case. Especially when IIUC, one can get the complex Jacobian or Wirtinger derivatives from such rules with just two backwards passes (see e.g. Complex Differentiation · Zygote).
Here’s a proposal for how to define complex chain rules:
frule and rrule assume f to be holomorphic and return the holomorphic derivative.Wirtinger wrapper type. If done so, then:frule will accept a pair (dz_dw,dz_dwb) as the corresponding input derivative and return a pair (df_dw,df_dwb) as the corresponding output derivative, andrrule will accept a pair (dv_df,dv_dfb) as the corresponding input derivative and return a pair (dv_dz,dv_dzb) as the corresponding output derivative.Throughout this post, a suffix b stands for “bar”, i.e. wb = \bar w.
For illustration, the chain rules for sin would be implemented as follows.
frule((_,dz_dw), ::typeof(sin), z::Number) = sin(z), cos(z)*dz_dw
rrule(::typeof(sin), z::Number) = sin(z), dv_dz->(NO_FIELDS, dv_df*cos(z))
frule((_,(dz_dw,dz_dwb)), ::typeof(sin), z::Wirtinger{<:Number}) =
sin(remove_wirtinger(z)), (Zero(), cos(remove_wirtinger(z)).*(dz_dw,dz_dwb))
rrule(::typeof(sin), z::Wirtinger{<:Number}) =
sin(remove_wirtinger(z), ((dv_df,dv_dfb),)->(NO_FIELDS, (dv_df*cos(remove_wirtinger(z)), dv_dfb*conj(cos(remove_wirtinger(z)))))
There are a couple of nice things about this approach:
Complex or Array{Complex}, i.e. no special numerical type needs to be introduced. This should avoid any compatibility issues with other libraries (most notably BLAS).f(a::Real,z::Complex) = real(a*z)
derivative(f, a::Real, z::Wirtinger) = (real(remove_wirtinger(z)), (a/2,a/2))
I suspect that the last point is more important to me than to most other people, and that there will be some debate about whether or not to make holomorphic derivatives the default; hence I would like to point out that my proposal of making holomorphic derivatives the default is based not only on my personal preferences but also on some technical arguments.
As an alternative to the above, we could define that complex arguments automatically imply Wirtinger derivatives, but this leads to the following issue: it makes perfectly sense to pass Wirtinger derivatives through \mathbb{R} \to \mathbb{R} functions (e.g. when computing the derivative of a function like \log(|z|)). Without the Wirtinger argument wrapper, every frule and rrule would thus have to be prepared to propagate either a single real derivative or a pair of Wirtinger derivatives, including all \mathbb{R} \to \mathbb{R} frule's and rrules. Making Wirtinger derivatives the default would hence require that all of ChainRules is aware of Wirtinger derivatives and handles them appropriately, which I think is not reasonable.
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?
Yeah, the PR skips abs and abs2 and uses Zygote’s: https://github.com/FluxML/Zygote.jl/blob/bf913a2a8ed616242e2f5378fbe598b289dd550a/src/lib/number.jl#L26-L30
I see why gradient(f,x) does the right thing for the cases 1) f holomorphic and 2) f(::Complex)::Real, but I don’t see why back(1),back(im) gives the Jacobian in the nonholomorphic case. What is the mathematical definition of pullback(f,x)?
The pullback of a function in Zygote is a function \mathcal{B_f}(v) = J_f^\dagger * v where v is a vector, J_f is the Jacobian of the function f and ^\dagger is the Hermitian conjugate.
Zygote will treat a struct with n fields as if it were a vector of length n, so the complex number 1 + 0im is treated like the vector [1 0] and 0 + 1im is treated like the vector [0, 1] in the above Jacobian vector product. The vector returned from this product is then converted into a struct
That is, in this case \mathcal{B}_f(1) = J_f^\dagger * [1, 0] and \mathcal{B}_f(i) = J_f ^\dagger * [0, 1]. The Jacobian here is 2x2, so this function returns a vector with two elements which are then interpreted as the real and imaginary parts of a complex number.
This is how Zygote deals with all structs, and it’s quite beautiful imo.
Thanks for your explanations, makes sense now.
This indeed seems to solve all issues very elegantly. The only thing I am still worried about is that holomorphic derivatives have to be implemented as
@scalar_rule f(z) conj(df(z))
It is too easy to forget the conj() and then waste hours figuring out why the resulting derivatives are not correct. I can see why doing it this way is consistent with interpreting complex derivatives as the “gradients” \tfrac{\partial u}{\partial x} + \iota \tfrac{\partial u}{\partial y}, but maybe it would nevertheless be easier to define complex derivatives as \tfrac{\partial u}{\partial x} - \iota \tfrac{\partial u}{\partial y} and avoid the conj() issue.
Yep, I couldn’t workout why it wasn’t working,
and I thought our defintion was copied directly from DiffRules (it might have been),
and Zygote has special handling for abs and abs2 in its DiffRules interactions so i figured that it would also need special handling for its abs and abs2 ChainRule interactions.
Do they? For scalar rules, both pushforwards and pullbacks can be computed by multiplying the sensitivity by the derivative of f. In Zygote’s approach to pullbacks, that derivative needs to be conjugated, but in forward mode it doesn’t. If that convention was adopted by ChainRules, you’d still want to implement the un-conjugated derivative, and then the @scalar_rule macro would conjugate in the rrule it creates. The rules themselves wouldn’t need to change. Unless I’m missing something.
You are right. Here’s a detailed write-up.
It seems that Zygote’s convention is that
# Not ChainRules.jl syntax.
# I merge the pullback call for notational convenience.
rrule(
z->u(z)+v(z)*im,
x+y*im,
dg_du+dg_dv*im
) -> dg_dx + dg_dy*im
represents the vector-matrix product
The analogous definition for frule would then be that
# Again not ChainRules.jl syntax.
# I rearranged arguments.
frule(
z->u(z)+v(z)*im,
x+y*im,
dx_dt+dy_dt*im
) -> du_dt + dv_dt*im
represents the matrix-vector product
If f(x+\iota y) = u(x,y) + \iota v(x,y) is holomorphic, then we want to implement these linear-algebra products as complex multiplication with \tfrac{df}{dz} = \tfrac{du}{dx} - \iota \tfrac{du}{dy}. We have
Thus rrule can indeed be implemented as multiplication with \overline{\tfrac{df}{dz}}, while frule can be implemented as multiplication with \tfrac{df}{dz}.
It looks like the Zygote- / “treat complex numbers as 2-vectors”-convention detailed above ticks all the right boxes:
The last point is one which has not really been addressed in the conversation so far, so I should elaborate a bit. Consider a function f : \mathbb{R}^n \to \mathbb{R}^m which internally uses some complex arithmetic. Zygote’s convention of treating structs like vectors ensures that the derivative of such a function is defined and agrees with what we would expect.
If there is agreement that this is what ChainRules.jl should implement, then I can give a hand with making sure that ChainRules.jl indeed follows this convention.
Would you mind trying to help me get this PR working: https://github.com/JuliaDiff/ChainRules.jl/pull/196 ?
I don’t understand Chainrules’ test suite very well and haven’t had the time lately to learn to fix the errors I’m getting.
I haven’t participated very actively in this thread, but it’s definitely great to see some momentum in this direction! I have written down some thoughts of my own in ChainRulesCore.jl#165. It doesn’t address C->C yet, but I would really appreciate some feedback from people here.
For reference, I opened a PR which rigorously defines how chain rules are to be implemented for complex functions.