vjp handles complex Jacobians well, but not jvp?

[xaviergonzalez]

Can anyone spot the bug in this jax code I've written? For some reason when I try to get the Jacobian using jvp, it eats my imaginary values, though I'm able to get vjp to work

import jax
import jax.numpy as jnp
from jax import vmap, jvp, vjp, jacrev, jacfwd
def fxn(z):
  return jnp.vdot(z,z).real
z0 = jnp.array([1j, 1j])
def get_jacobian_fwd(f, p):
  """
  Return the Jacobian of a complex to real function f
  """
  _, J = vmap(lambda t: jvp(f, (p,), (t,)))(jnp.eye(len(p), dtype=jnp.complex64))
  return J
def get_jacobian_rev(f,p):
  _, f_vjp = vjp(f,p)
  return f_vjp(1.)

get_jacobian_fwd(fxn, z0) gives: Array([0., 0.], dtype=float32)
get_jacobian_rev(fxn, z0) gives: (Array([0.-2.j, 0.-2.j], dtype=complex64),)

[mattjj]

Thanks for the question!

Since the primal outputs of your computation are always real (because of the .real on the end), the perturbations to them (ie tangents) will be too. In other words, if you have a function $f : \mathbb{C} \to \mathbb{R}$, then $\partial f : \mathbb{C} \to \mathbb{C} \to \mathbb{R}$. And the JVP evaluates $(x, , \dot x) \mapsto (f(x), , \partial f(x) , \dot x) \in \mathbb{R} \times \mathbb{R}$.

Maybe the issue is that you're trying to compute the Jacobian matrix in two ways, and get the same answer? But the full Jacobian of a $\mathbb{C}^2 \to \mathbb{R}$ function must have 4 real coefficients. In particular, get_jacobian_fwd is currently just trying the vectors (1, 0) and (0, 1), but it needs to also try (j, 0) and (0, j). That is, try replacing the jnp.eye with jnp.array([[1., 0.], [0., 1.], [1j, 0.], [0., 1j]]). Then you'll get something isomorphic to the same answer as you get in the VJP case.

What do you think?