Autograd for Complex Numbers

**What notion of complex derivative does PyTorch use?**

PyTorch follows `JAX's <https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#Complex-numbers-and-differentiation>`_

convention for autograd for Complex Numbers.

Suppose we have a function :math:`F: ℂ → ℂ` which we can decompose into functions u and v

which compute the real and imaginary parts of the function:

.. code::

where :math:`1j` is a unit imaginary number.

We define the JVP for function :math:`F` at :math:`(x, y)` applied to a tangent

vector :math:`c+dj \in C` as:

.. code::

where

This is similar to the definition of the JVP for a function defined from :math:`R^2 → R^2`, and the multiplication

with :math:`[1, 1j]^T` is used to identify the result as a complex number.

We define the VJP of :math:`F` at :math:`(x, y)` for a cotangent vector :math:`c+dj \in C` as:

.. code::

In PyTorch, the VJP is mostly what we care about, as it is the computation performed when we do backward

mode automatic differentiation. Notice that d and :math:`1j` are negated in the formula above.

**Why is there a negative sign in the formula above?**

For a function F: V → W, where are V and W are vector spaces. The output of

the Vector-Jacobian Product :math:`VJP : V → (W^* → V^*)` is a linear map

from :math:`W^* → V^*` (explained in `Chapter 4 of Dougal Maclaurin’s thesis <https://dougalmaclaurin.com/phd-thesis.pdf>`_).

The negative signs in the above `VJP` computation are due to conjugation. :math:`c-dj`

is the covector in dual space of :math:`C^` (:math:`\in ℂ^*`) corresponding to the

cotangent vector :math:`c+dj`, and the multiplication by :math:`[1, -1j]`, whose net effect is

to get a conjugate of the complex number we would have obtained by multiplcation with :math:`[1, 1j]` instead,

is used to get the result in :math:`ℂ` since the final result of reverse-mode differentiation of a function

is a covector belonging to :math:`ℂ^*` (explained in

`Chapter 4 of Dougal Maclaurin’s thesis <https://dougalmaclaurin.com/phd-thesis.pdf>`_).

**What happens if I call backward() on a complex scalar?**

The gradient for a complex function is computed assuming the input function is a holomorphic function.

This is because for general :math:`ℂ → ℂ` functions, the Jacobian has 4 real-valued degrees of freedom

(as in the `2x2` Jacobian matrix above), so we can’t hope to represent all of them with in a complex number.

However, for holomorphic functions, the gradient can be fully represented with complex numbers due to the

Cauchy-Riemann equations that ensure that `2x2` Jacobians have the special form of a scale-and-rotate

matrix in the complex plane, i.e. the action of a single complex number under multiplication. And so, we can

obtain that gradient using backward which is just a call to `vjp` with covector `1.0`.

The net effect of this assumption is that the partial derivatives of the imaginary part of the function

(:math:`v(x, y)` above) are discarded for :func:`torch.autograd.backward` on a complex scalar

(e.g., this is equivalent to dropping the imaginary part of the loss before performing a backwards).

For any other desired behavior, you can specify the covector `grad_output` in :func:`torch.autograd.backward` call accordingly.

**How are the JVP and VJP defined for cross-domain functions?**

Based on formulas above and the behavior we expect to see (going from :math:`ℂ → ℝ^2 → ℂ` should be an identity),

we use the following formula for cross-domain functions.

The JVP and VJP for a :math:`f1: ℂ → ℝ^2` are defined as:

.. code::

The JVP and VJP for a :math:`f1: ℝ^2 → ℂ` are defined as:

.. code::