Repository: google/jax

Issue:
Inconsistent convention for complex derivatives
The convention for the handling of complex derivatives in JAX seems to be mixed. Consider the following two implementations of the abolute value function.
```py
import jax.numpy as np
from jax import grad

def f(x):
    return np.abs(x)

def g(x):
    return np.sqrt(np.real(x)**2 + np.imag(x)**2)

df = grad(f)
dg = grad(g)

print('gradient by df: ',df(1.0+0.5j))
print('gradient by dg: ',dg(1.0+0.5j))
```
The respective results differ by complex conjugation:
```
gradient by df:  (0.8944272-0.4472136j)
gradient by dg:  (0.8944272+0.4472136j)
```


Discussion:
Thanks for catching this! To check against our [Autograd](https://github.com/hips/autograd) implementation, I changed the two import lines at the top to

```python
import autograd.numpy as np
from autograd import grad
```

And got this:

```
('gradient by df: ', (0.8944271909999159-0.4472135954999579j))
('gradient by dg: ', (0.8944271909999159-0.4472135954999579j))
```

So we've got a bug here, but at least we got it right when we wrote this in Autograd, and can perhaps compare to Autograd's implementation to figure out where we're going wrong.
The same issue appears if we get rid of the `np.sqrt` and evaluate `np.abs(x) ** 2`.

This may be a discrepancy: for `imag`, Autograd's [forward-mode](https://github.com/HIPS/autograd/blob/387c373115ddd54cff2c8ba6a9fc619f28639cfb/autograd/numpy/numpy_jvps.py#L119) and [reverse-mode](https://github.com/HIPS/autograd/blob/387c373115ddd54cff2c8ba6a9fc619f28639cfb/autograd/numpy/numpy_vjps.py#L126) rules both multiply by `-1j`, whereas in [JAX's rules](https://github.com/google/jax/blob/527fe14838eadc48fede68bd42ce4364e269fed3/jax/lax/lax.py#L1611) we don't and instead just treat it as a linear projection on R^2.

That's a lead on what's going wrong. Separately I want to figure out why our tests didn't catch this.
I think Autograd's JVP effectively just takes the imaginary part (because the Autograd JVP effectively evaluates `real(-1j * g)` and so if `g = a + bj` then `-1j * (a + bj) = b - aj`), which matches [JAX's JVP](https://github.com/google/jax/blob/master/jax/lax/lax.py#L1611) because `ad.deflinear` just means apply the primitive to the tangent.

But Autograd's VJP rule, which corresponds to composing JAX's JVP rule with its transpose rule, introduces a negative sign: it evaluates `-1j * g`, whereas [JAX's transpose rule doesn't include a negative sign](https://github.com/google/jax/blob/527fe14838eadc48fede68bd42ce4364e269fed3/jax/lax/lax.py#L1611).

By replacing JAX's transpose rule with this, all our tests pass _and_ the code in this thread works as expected:

```python
ad.deflinear(imag_p, lambda t: [complex(onp.zeros((), _dtype(t)), neg(t))])
```

But I still don't have a clear view of what's going on. The transpose rules are meant to be transposes rather than adjoints, and so I'd thought we could model `imag` as the action of the matrix `[0 1]` on R^2...
Finally getting back to this! And I paged back in the [complex number convention](https://github.com/HIPS/autograd/blob/master/docs/tutorial.md#complex-numbers).

We define JVPs and VJPs for a function f : C -> C in terms of this matrix as a linear map R^2 -> R^2:

![image](https://user-images.githubusercontent.com/1458824/60742793-9e4c8f00-9f23-11e9-9aa4-0bc2c1cf2839.png)

where f(x, y) = u(x, y) + i * v(x, y). (We do this because if f is holomorphic then this matrix recovers the complex derivative, i.e. it's the real matrix representation of a single complex number, and because it looks just like the Jacobian of f up to flipping the signs of the off-diagonals, i.e. up to complex conjugation of the input and output.)

In particular, to model the `imag` function we have f(x, y) = (y, 0) and so the matrix above has the form

![image](https://user-images.githubusercontent.com/1458824/60742495-70b31600-9f22-11e9-9afa-da5c20187705.png)

That's where the negative comes from!

As to why the tests didn't catch it, it seems [we left some complex derivative tests commented out](https://github.com/google/jax/blob/527fe14838eadc48fede68bd42ce4364e269fed3/tests/lax_test.py#L1526-L1527) and then forgot about them! Oops.

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1608,12 +1608,12 @@ def _brcast_to(x, shape):
 ad.deflinear(real_p, lambda t: [complex(t, onp.zeros((), _dtype(t)))])
 
 imag_p = unop(_complex_basetype, _complex, 'imag')
-ad.deflinear(imag_p, lambda t: [complex(onp.zeros((), _dtype(t)), t)])
+ad.defjvp(imag_p, lambda g, _: real(mul(_const(g, -1j), g)))
 
 _complex_dtype = lambda dtype, *args: (onp.zeros((), dtype) + onp.zeros((), onp.complex64)).dtype
 complex_p = binop(_complex_dtype, [_complex_elem_types, _complex_elem_types],
                   'complex')
-ad.deflinear(complex_p, lambda t: [real(t), imag(t)])
+ad.deflinear(complex_p, lambda t: [real(t), imag(neg(t))])
 
 conj_p = unop(_complex_dtype, _float | _complex, 'conj')
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1523,10 +1523,10 @@ def grad_test_spec(op, nargs, order, rng, dtypes, name=None, tol=None):
 
     grad_test_spec(lax.real, nargs=1, order=2, rng=jtu.rand_default(),
                    dtypes=[onp.complex64]),
-    # grad_test_spec(lax.imag, nargs=1, order=2, rng=jtu.rand_default(),
-    #                dtypes=[onp.complex64]),  # TODO(mattjj): enable
-    # grad_test_spec(lax.complex, nargs=2, order=2, rng=jtu.rand_default(),
-    #                dtypes=[onp.float32]),  # TODO(mattjj): enable
+    grad_test_spec(lax.imag, nargs=1, order=2, rng=jtu.rand_default(),
+                   dtypes=[onp.complex64]),
+    grad_test_spec(lax.complex, nargs=2, order=2, rng=jtu.rand_default(),
+                   dtypes=[onp.float32]),
     grad_test_spec(lax.conj, nargs=1, order=2, rng=jtu.rand_default(),
                    dtypes=[onp.float32, onp.complex64]),
     grad_test_spec(lax.abs, nargs=1, order=2, rng=jtu.rand_positive(),