vjp for a complex function with real inputs gives wrong results

[fishjojo]

I'd like to compute the derivative of the function exp(-iR⋅G) (where R and G are real) wrt R, which is simply -iG exp(-iR⋅G). The code is as follows:

import numpy
import jax
from jax import numpy as jnp

def get_SI(R,G):
    return jnp.exp(-1j * jnp.dot(R,G.T))

natm = 2
ng = 10
R = jnp.ones((natm,3))
G = numpy.random.rand(ng,3)

jac_fwd = jax.jacfwd(get_SI)(R,G)
print(jac_fwd.shape, jac_fwd.dtype)

#R = jnp.asarray(R, dtype=jnp.complex64)
_, func_vjp = jax.vjp(get_SI, R, G)
SI_bar = jnp.eye(natm*ng, dtype=jnp.complex64).reshape(-1,natm,ng)
jac_bwd = jax.vmap(func_vjp)(SI_bar)[0].reshape(natm,ng,natm,3)
print(jac_bwd.shape, jac_bwd.dtype)
print(abs(jac_bwd - jac_fwd).max())

print(jax.make_jaxpr(func_vjp)(SI_bar[0]))

and the output is

(2, 10, 2, 3) complex64
(2, 10, 2, 3) float32
0.31839034
{ lambda a b c ; d.
  let e = mul d a
      f = mul e (-0-1j)
      g = convert_element_type[ new_dtype=float32
                                weak_type=False ] f
      h = dot_general[ dimension_numbers=(((0,), (0,)), ((), ()))
                       precision=None
                       preferred_element_type=None ] g b
      i = transpose[ permutation=(1, 0) ] h
      j = transpose[ permutation=(1, 0) ] i
      k = dot_general[ dimension_numbers=(((1,), (1,)), ((), ()))
                       precision=None
                       preferred_element_type=None ] g c
  in (k, j) }

You can see that the forward mode gives correct answer, but the reverse mode (by vjp) has a type conversion which leads to real derivatives. If I convert the input to complex numbers, the reverse mode works right. Is this a bug or am I doing anything wrong?
Thanks.

[mattjj]

Thanks for the question! If you haven't taken a look already, I suggest reading the autodiff cookbook section on complex numbers and differentiation.

In general, the VJP of an R->C function will be a C->R function. That is, it'll always produce real results, and so the behavior you described is working as intended.

Fundamentally this comes down to the difference between JVPs/VJPs and the coefficients of a Jacobian matrix. I think you were expecting the (scalar) Jacobian, since indeed that's what the expression -iG exp(-iR⋅G) evaluates to. That's a reasonable thing to want! To get a Jacobian from autodiff, we have to figure out how to get it from JVPs and VJPs.

As you observed, if f : R->C then one way to get its Jacobian at x is to use jvp(f, (x,), (1.,)), since that'll be the complex number we're after. One way to think about this is that if we identify f with g: R->R^2, then g has a Jacobian of shape (2, 1), and we can use a single JVP to reveal both its real coefficients at once (or correspondingly its one complex coefficient). But we can't do the same with a single VJP.

Luckily, since exp is holomorphic, we can just lift it to a C->C function and use a single VJP (or JVP) to get its Jacobian represented as a single complex number. That amounts to doing exactly what you said: just pass a complex input. That seems like the best approach in this case.

What do you think?

[fishjojo]

Thanks for the detailed explanation. I guess there is no easy way to get the complex gradient for R->C functions with VJP. For simple functions like f=cx, one might be able to add an option to decide whether to transpose the type conversion of x, but it may be wrong for general cases. I would be happy to know if there is any news for this kind of problems.

[PhilipVinc]

@fishjojo over in NetKet we had a lot of issues with that, and we ended up wrapping jax.vjp into our own nk.jax.vjp to automatically handle such cases, that are very common in quantum mechanics.
We now use nk.jax.vjp as a drop-in replacement to jax.vjp in our code and never worry about whever our function is R->R, R->C, C->C and what the vector is.

The code is fairly straightfoward and you can have a look here

[fishjojo]

Thanks @PhilipVinc, this is very helpful.