Wirtinger derivative using JVP in jax

[kcdodd]

I am working to better understand how automatic differentiation relates to more theoretical complex analysis of non-holomorphic functions, and looking for feedback on the following way of thinking or explaining it.

As a theoretical beginning, defining the JVP ("Jacobian vector prduct") as the directional derivative $J(v)$ ,

$$\lim_{h \to 0} \frac{\left|f(z_0 + h v) - f(z_0) - h J(v)\right|}{h} = 0.$$

The JVP doesn't always give what might be expected because the value of $J(v)$ is "like" the value of $f(z_0)$. For example $f(z) = |z|^2$ the value of $f$ is real, so the JVP has to be real regardless of the tangent $v$. Treating the mapping like a matrix-vector product is troubled for complex number: e.g. for complex $c$, $J(c v) \ne c J(v)$. Complex numbers are like 2D transformations that can scale and rotate coordinates in the complex plane, in addition to also being coordinates themselves. For example, with the matrix representation of a complex number,

$$z = a + b i \to a \begin{pmatrix} 1 & 0 \\\ 0 & 1 \end{pmatrix} + b \begin{pmatrix} 0 & -1 \\\ 1 & 0 \end{pmatrix} = \begin{pmatrix} a & -b \\\ b & a \end{pmatrix}.$$

But complex numbers cannot represent things like conjugation (reflection), so the Jacobian does not always have a representation as a purely complex-valued matrix. So starting instead with the 2D (real) coordinate representation, and a general real-valued 2x2 matrix (transformation),

$$A = \begin{pmatrix} A_{11} & A_{12} \\\ A_{21} & A_{22} \end{pmatrix},$$

the question becomes how to reconstruct the result of the transformation using JVP. The first response would be to separately compute the real and imaginary directions, and basically yes, but more exactly how to combine the results. In the case of $f(z) = |z|^2$ for example, the JVP ends up aligned with the real axis no matter how the tangent is oriented. The following approach is to decompose $A$ as a linear combination of elementary transformations,

$$A = \alpha_1 \sigma_1 + \alpha_2 \sigma_2 + \alpha_3\sigma_3 + \alpha_4\sigma_4.$$

The first two matrices are to be identified with orientation preserving scale and rotation ($1$ and $i$), while the second two involve reflection/conjugation,

$$\begin{align*} &z \to z, & &z \to i z, & &z \to \bar{z}, & &z \to i\bar{z}, \\\ \sigma_1 =& \begin{pmatrix} 1 & 0 \\\ 0 & 1 \end{pmatrix}, &\quad \sigma_2 =& \begin{pmatrix} 0 & -1 \\\ 1 & 0 \end{pmatrix}, &\quad \sigma_3 =& \begin{pmatrix} 1 & 0 \\\ 0 & -1 \end{pmatrix}. &\quad \sigma_4 =& \begin{pmatrix} 0 & 1 \\\ 1 & 0 \end{pmatrix}, \\\ \alpha_1 =& \frac{A_{11} + A_{22}}{2}, &\quad \alpha_2 =& \frac{A_{21} - A_{12}}{2}, &\quad \alpha_3 =& \frac{A_{11} - A_{22}}{2}, &\quad \alpha_4 =& \frac{A_{21} + A_{12}}{2}. \end{align*}$$

They also have these identities,

$$\begin{align*} \sigma_2^2 &= -\sigma_1, \quad\sigma_2^{-1} = \sigma_2^T = -\sigma_2, \\\ \sigma_3^2 &= \sigma_4^2 = \sigma_1, \\\ \sigma_2 &= \sigma_4 \sigma_3 = -\sigma_3 \sigma_4, \\\ \sigma_3 &= \sigma_4 \sigma_2 = -\sigma_2 \sigma_4, \\\ \sigma_4 &= \sigma_2 \sigma_3 = - \sigma_3 \sigma_2. \\\ \end{align*}$$

Since $\sigma_4 = \sigma_2 \sigma_3$ the decomposition can be expressed with only $\sigma_1$, $\sigma_2$, and $\sigma_3$, which gives a nice picture of how to connect the action on complex values to a general real-valued matrix,

$$\begin{align*} A &= \alpha_1 \sigma_1 + \alpha_2 \sigma_2 + \left(\alpha_3 + \alpha_4\sigma_2\right)\sigma_3, \\\ z &\to (\alpha_1 + i\alpha_2)z + \left(\alpha_3 + i\alpha_4\right)\bar{z}. \end{align*}$$

Now suppose a second matrix is defined by a similarity transformation using $\sigma_2$ (aka $i$), which basically rotates by $\pi/2$ before the matrix and then rotates back,

$$\begin{align*} A^\prime &= \sigma_2^{-1} A \sigma_2 = -\sigma_2 A \sigma_2 \\\ &= \alpha_1 \sigma_1 + \alpha_2 \sigma_2 - \left(\alpha_3 + \alpha_4\sigma_2\right)\sigma_3. \end{align*}$$

The reflected part is reversed because it rotated in the opposite direction. The sum $A + A^\prime$ will keep only the un-reflected part, while the difference $A - A^\prime$ will keep only the reflected part. In the Wirtinger derivatives, $z$ and $\bar{z}$ are treated as though they were independent coordinates of an effective function $f(z, \bar{z})$. Identifying these with separate actions on a complex value $z$ and its conjugate $\bar{z}$, an equivalent of the Wirtinger derivatives can be defined for the JVP. Substituting the matrix representations, $A \to J(v)$ and $A^\prime \to - i J(i v)$,

$$J_z(v) = \frac{J(v) - i J(i v)}{2},$$ $$J_{\bar{z}}(v) = \frac{J(v) + i J(i v)}{2}.$$

Something to consider when interpreting the conjugate derivative is that it is not quite the same thing as the derivative along the tangent's conjugate direction $\bar{v}$, but only the part where the original function depended on $\bar{z}$. For a holomorphic function $J_{\bar{z}}(v)$ is zero, while using the conjugate tangent in $J_{z}(\bar{v})$ is just taking the derivative with respect to $z$ along a different direction.

An example implementation using jax is shown in the code below.

import jax
from jax import numpy as jnp
from jax.tree_util import tree_map

def jvp_wirtinger(f, primals, tangents):
  assert type(primals) is type(tangents)

  if not isinstance(primals, tuple):
    primals = (primals,)
    tangents = (tangents,)

  primals_out, f_jvp = jax.linearize(f, *primals)
  # J(v)
  Ja = f_jvp(*tangents)
  # - i J(i v)
  Jb = tree_map(lambda v: -1j*v, f_jvp(*tree_map(lambda v: 1j*v, tangents)))

  Jz = tree_map(lambda a, b: 0.5*(a + b), Ja, Jb)
  Jzconj = tree_map(lambda a, b: 0.5*(a - b), Ja, Jb)
  return Jz, Jzconj