Repository: google/jax

Issue:
Error in backward-mode differentiation of ℝ → ℂ function
Hi,
First of all I'd like to thank you for the wonderful work you put forward with Jax. 
We recently started using it in [netket](https://github.com/netket/netket), a package for solving quantum-many body problems with neural-networks, and have been very pleasantly surprised by the performance and ease of use, so much that we rewrote most of our code to take advantage of it!

However, when dealing with wavefunctions we have to work a lot with complex functions. It is within this context that I noticed a silent error arises when trying to use reverse-mode over a ℝ → ℂ function (so, `grad` and `jacrev`).

Consider the following case:
```python
import jax
w=(jax.np.array([1.0]), jax.np.array([0.5]))

jax.grad(lambda w: w[0].sum() + 1j* w[1].sum(), holomorphic=True)(w)
(DeviceArray([1.], dtype=float32), DeviceArray([0.], dtype=float32))

 jax.jacrev(lambda w: w[0].sum() + 1j* w[1].sum(), holomorphic=True)(w)
(DeviceArray([1.], dtype=float32), DeviceArray([0.], dtype=float32))
```
- The gradient is wrong, it should be complex-valued. I have the feeling you are probably competing the right gradient but you cast it at the end to the original type, resulting in this error.

As a temporary fix, I figured out I can just promote types to be complex and compute the gradient by saying that the function is holomorphic, and re-scaling the complex part, though it would be nice if this could be fixed.
```python 
wc=(jax.np.array([1.0], dtype=jax.np.complex64), jax.np.array([0.5], dtype=jax.np.complex64))

jax.jacrev(lambda w: w[0].sum() + 1j* w[1].sum(), holomorphic=True)(wc)
(DeviceArray([1.-0.j], dtype=complex64), DeviceArray([0.+1.j], dtype=complex64))
```

Slightly related,  I understand why you ask the user to mark the function to be holomorphic once you hit a complex result, but in this case it's a bit weird: ℝ → ℂ functions are always holomorphic. I don't know if anything can be done about it.

cc @gcarleo


Discussion:
Yes, thank you, we are really enjoying the clarity of Jax's APIs! It makes working with complex-valued quantum wave functions much easier and requires a lot less boilerplate than with other approaches. 

Something related to what @PhilipVinc was saying (in the same spirit of mixed real complex typing). There is this slight issue with vjp for complex vectors and real functions and vice versa:  
 
```python
import jax
import jax.numpy as jnp


f = jnp.sin

x = jnp.ones(3)
v = jnp.ones(3)

# R->R Works
val, f_jvp = jax.vjp(f, x)

print(f_jvp(v))


# X Real and V complex Works only if casting x to complex
x = jnp.ones(3, dtype=jnp.complex64)
v = jnp.ones(3) + 1.0j * jnp.ones(3)

val, f_jvp = jax.vjp(f, x)
print(f_jvp(v))


# X Complex and V real Works only if casting v to complex
x = jnp.ones(3) + 1.0j * jnp.ones(3)
v = jnp.ones(3, dtype=jnp.complex64)

val, f_jvp = jax.vjp(f, x)
print(f_jvp(v))
```

The second case can be implemented calling two jvp with real x, instead of casting x to complex (most likely that is going to be also more efficient than casting). I was wondering if you have plans to implement these mixed-typing cases directly into the AD engine or if we should handle them separately in our code. 

Thank you! 
Broadly speaking it's perfectly reasonable to implement the complex AD semantics you need in your own code: everything is pretty easy to construct out of unambiguous R -> R derivatives. JAX [makes particular choices in `grad`](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#Complex-numbers-and-differentiation) that are different from some other libraries: we treat a R -> C or C -> C function as if it were a function onto the reals (i.e., we throw out the complex component). In the special case where f is holomorphic C -> C, that definition lines up with the complex derivative; in other cases, it doesn't.
I haven't paged this issue in yet, but I wonder if the comment on #610 might be useful.
Yes, that looks very relevant and perhaps worth making an FAQ!
Thank you @jekbradbury  and @mattjj for the quick feedback! 

Indeed I agree that there is absolutely no ambiguity in the vjp APIs, and the conventions taken in JAX for complex numbers are clear. 
I was just remarking that it might be quite convenient for users if vjp (and jvp) had a mixed-type usage consistent with matmul. 

For example, a matmul with complex mat and real vector, just works in numpy: 

```python
import jax
import jax.numpy as jnp

# Real matrix
J = jnp.ones((3, 3))

# Complex vector
v = (1.0 + 1.0j) * jnp.ones(3)

# Mixed typing works without casting
out = jnp.matmul(J, v)
```

however an equivalent vjp wouldn't work unless one manually performs a tree flattening and unflattening

```python
# Function of real parameters
x = jnp.ones(3)
f_jvp = jax.vjp(jnp.sin, x)[1]

## Jacobian-vector product doesn't work without casting
# out = f_jvp(v)

## This **would** work if addition and scalar multiplication were overloaded for PyTrees
# out = f_jvp(v.real) + 1.0j * f_jvp(v.imag)

# This works (but it might be non-trivial for some users)
from jax.tree_util import tree_map, tree_flatten, tree_unflatten

res, td = tree_flatten(f_jvp(v.real))
ims, td = tree_flatten(f_jvp(v.imag))

out = tree_unflatten(td, [re + 1.0j * im for re, im in zip(res, ims)])
```

In the example above, casting x to complex might also be an option, however for non-linear functions I suspect that computing a jvp for complex parameters is more than twice as expensive as calling f_jvp on real parameters? 

Hmm, well in your example you have a complex vector input, and that works with `jacfwd` / `jacrev` as well because you're clearly modeling a C->C function:

```python
import jax
import jax.numpy as jnp
from jax import jacfwd, jacrev

# Real matrix
J = jnp.ones((3, 3))

# Complex vector
v = (1.0 + 1.0j) * jnp.ones(3)

def f(v):
  return jnp.matmul(J, v)

print(jacrev(f, holomorphic=True)(v))  # notice complex v input
```

```
[[1.+0.j 1.+0.j 1.+0.j]
 [1.+0.j 1.+0.j 1.+0.j]
 [1.+0.j 1.+0.j 1.+0.j]]
```

Calling it at `jnp.ones_like(v)` works the same way, but not `jnp.ones(v.shape)`.

I think the fundamental issue here is JAX can't know that you think of `f` as a `C->C` function unless you pass in a complex input, and therefore can't know to construct an input basis with complex dtype, because Python is polymorphic and we only ascribe an input type to `f` right when we execute it using the inputs you provide.

We could:
1. raise an error if you pass `holomorphic=True` but don't provide a complex dtype input, or
2. automatically cast inputs to complex dtype when `holomorphic=True`.

I think you might be requesting the latter.

WDYT?

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -150,7 +150,7 @@ def f_jitted(*args, **kwargs):
     else:
       dyn_args = args
     args_flat, in_tree = tree_flatten((dyn_args, kwargs))
-    _check_args(args_flat)
+    for arg in args_flat: _check_arg(arg)
     flat_fun, out_tree = flatten_fun(f, in_tree)
     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
                        name=flat_fun.__name__)
@@ -370,7 +370,7 @@ def grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
       first element is considered the output of the mathematical function to be
       differentiated and the second element is auxiliary data. Default False.
     holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be
-      holomorphic. Default False.
+      holomorphic. If True, inputs and outputs must be complex. Default False.
 
   Returns:
     A function with the same arguments as ``fun``, that evaluates the gradient
@@ -424,7 +424,7 @@ def value_and_grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
      first element is considered the output of the mathematical function to be
      differentiated and the second element is auxiliary data. Default False.
     holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be
-      holomorphic. Default False.
+      holomorphic. If True, inputs and outputs must be complex. Default False.
 
   Returns:
     A function with the same arguments as ``fun`` that evaluates both ``fun``
@@ -454,17 +454,14 @@ def value_and_grad_f(*args, **kwargs):
 
     f = lu.wrap_init(fun, kwargs)
     f_partial, dyn_args = argnums_partial(f, argnums, args)
+    tree_map(partial(_check_input_dtype_grad, holomorphic), dyn_args)
     if not has_aux:
       ans, vjp_py = _vjp(f_partial, *dyn_args)
     else:
       ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)
     _check_scalar(ans)
     dtype = dtypes.result_type(ans)
-    if not (holomorphic or dtypes.issubdtype(dtype, onp.floating)):
-      msg = ("Gradient only defined for real-output functions (with dtype that "
-             "is a subdtype of np.floating), but got dtype {}. For holomorphic "
-             "differentiation, pass holomorphic=True.")
-      raise TypeError(msg.format(dtype))
+    tree_map(partial(_check_output_dtype_grad, holomorphic), ans)
     g = vjp_py(onp.ones((), dtype=dtype))
     g = g[0] if isinstance(argnums, int) else g
     if not has_aux:
@@ -487,6 +484,38 @@ def _check_scalar(x):
     else:
       raise TypeError(msg("had abstract value {}".format(aval)))
 
+def _check_input_dtype_revderiv(name, holomorphic, x):
+  _check_arg(x)
+  aval = core.get_aval(x)
+  if holomorphic:
+    if not dtypes.issubdtype(aval.dtype, onp.complexfloating):
+      msg = (f"{name} with holomorphic=True requires inputs with complex dtype, "
+             f"but got {aval.dtype.name}.")
+      raise TypeError(msg)
+  elif not (dtypes.issubdtype(aval.dtype, onp.floating) or
+            dtypes.issubdtype(aval.dtype, onp.complexfloating)):
+    msg = (f"{name} requires real- or complex-valued inputs (input dtype that "
+           "is a sub-dtype of np.floating or np.complexfloating), "
+           f"but got {aval.dtype.name}. ")
+    raise TypeError(msg)
+_check_input_dtype_grad = partial(_check_input_dtype_revderiv, "grad")
+
+def _check_output_dtype_revderiv(name, holomorphic, x):
+  aval = core.get_aval(x)
+  if holomorphic:
+    if not dtypes.issubdtype(aval.dtype, onp.complexfloating):
+      msg = (f"{name} with holomorphic=True requires outputs with complex dtype, "
+             f"but got {aval.dtype.name}.")
+      raise TypeError(msg)
+  elif not dtypes.issubdtype(aval.dtype, onp.floating):
+    msg = (f"{name} requires real-valued outputs (output dtype that is "
+           f"a sub-dtype of np.floating), but got {aval.dtype.name}. "
+           "For holomorphic differentiation, pass holomorphic=True. "
+           "For differentiation of non-holomorphic functions involving complex "
+           "outputs, use jax.vjp directly.")
+    raise TypeError(msg)
+_check_output_dtype_grad = partial(_check_output_dtype_revderiv, "grad")
+
 
 def jacfwd(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
            holomorphic: bool = False) -> Callable:
@@ -521,21 +550,39 @@ def jacfwd(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
   def jacfun(*args, **kwargs):
     f = lu.wrap_init(fun, kwargs)
     f_partial, dyn_args = argnums_partial(f, argnums, args)
-    holomorphic or tree_map(_check_real_input_jacfwd, dyn_args)
+    tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args)
     pushfwd = partial(_jvp, f_partial, dyn_args)
     y, jac = vmap(pushfwd, out_axes=(None, batching.last))(_std_basis(dyn_args))
+    tree_map(partial(_check_output_dtype_jacfwd, holomorphic), y)
     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
     return tree_map(partial(_unravel_array_into_pytree, example_args, -1), jac)
 
   return jacfun
 
-def _check_real_input_jacfwd(x):
+def _check_input_dtype_jacfwd(holomorphic, x):
+  _check_arg(x)
   aval = core.get_aval(x)
-  if not dtypes.issubdtype(aval.dtype, onp.floating):
-    msg = ("jacfwd only defined for functions with input dtypes that are "
-           "sub-dtypes of `np.floating` (i.e. that model real values), but "
-           "got {}. For holomorphic differentiation, pass holomorphic=True.")
-    raise TypeError(msg.format(aval.dtype.name))
+  if holomorphic:
+    if not (dtypes.issubdtype(aval.dtype, onp.complexfloating) and
+            not dtypes.issubdtype(aval.dtype, onp.floating)):
+      msg = ("jacfwd with holomorphic=True requires inputs with complex dtype, "
+             f"but got {aval.dtype.name}.")
+      raise TypeError(msg)
+  elif not dtypes.issubdtype(aval.dtype, onp.floating):
+    msg = ("jacfwd requires real-valued inputs (input dtype that is "
+           f"a sub-dtype of np.floating), but got {aval.dtype.name}. "
+           "For holomorphic differentiation, pass holomorphic=True. "
+           "For differentiation of non-holomorphic functions involving complex "
+           "inputs, use jax.jvp directly.")
+    raise TypeError(msg)
+
+def _check_output_dtype_jacfwd(holomorphic, x):
+  aval = core.get_aval(x)
+  if holomorphic:
+    if not dtypes.issubdtype(aval.dtype, onp.complexfloating):
+      msg = ("jacfwd with holomorphic=True requires outputs with complex dtype, "
+             f"but got {aval.dtype.name}.")
+      raise TypeError(msg)
 
 
 def jacrev(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
@@ -571,8 +618,9 @@ def jacrev(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
   def jacfun(*args, **kwargs):
     f = lu.wrap_init(fun, kwargs)
     f_partial, dyn_args = argnums_partial(f, argnums, args)
+    tree_map(partial(_check_input_dtype_jacrev, holomorphic), dyn_args)
     y, pullback = _vjp(f_partial, *dyn_args)
-    holomorphic or tree_map(_check_real_output_jacrev, y)
+    tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
     jac = vmap(pullback)(_std_basis(y))
     jac = jac[0] if isinstance(argnums, int) else jac
     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
@@ -582,13 +630,8 @@ def jacfun(*args, **kwargs):
   return jacfun
 jacobian = jacrev
 
-def _check_real_output_jacrev(x):
-  aval = core.get_aval(x)
-  if not dtypes.issubdtype(aval.dtype, onp.floating):
-    msg = ("jacrev only defined for functions with output dtypes that are "
-           "sub-dtypes of `np.floating` (i.e. that model real values), but "
-           "got {}. For holomorphic differentiation, pass holomorphic=True.")
-    raise TypeError(msg.format(aval.dtype.name))
+_check_input_dtype_jacrev = partial(_check_input_dtype_revderiv, "jacrev")
+_check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev")
 
 
 def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
@@ -1070,7 +1113,7 @@ def f_pmapped(*args, **kwargs):
     assert all(axis in (0, None) for axis in in_axes_flat), \
         "pmap currently only supports mapping over the leading axis"
     local_axis_size = _mapped_axis_size(in_tree, args, in_axes_flat, "pmap")
-    _check_args(args)
+    for arg in args: _check_arg(arg)
     flat_fun, out_tree = flatten_fun(f, in_tree)
     out = pxla.xla_pmap(
         flat_fun,
@@ -1114,7 +1157,7 @@ def f_pmapped(*args, **kwargs):
         "soft_pmap currently only supports mapping over the leading axis"
     mapped_invars = tuple(axis is not None for axis in in_axes_flat)
     axis_size = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "soft_pmap")
-    _check_args(args_flat)
+    for arg in args_flat: _check_arg(arg)
     flat_fun, out_tree = flatten_fun(f, in_tree)
 
     chunk_size, leftover = divmod(axis_size, pxla.unmapped_device_count(backend))
@@ -1489,7 +1532,7 @@ def _vjp(fun: lu.WrappedFun, *primals, **kwargs):
   has_aux = kwargs.pop('has_aux', False)
   assert not kwargs
   primals_flat, in_tree = tree_flatten(primals)
-  _check_args(primals_flat)
+  for arg in primals_flat: _check_arg(arg)
   tree_map(_check_inexact_input_vjp, primals)
   if not has_aux:
     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
@@ -1618,11 +1661,10 @@ def device_get(x):
   return tree_map(_device_get, x)
 
 
-def _check_args(args):
-  for arg in args:
-    if not (isinstance(arg, core.Tracer) or _valid_jaxtype(arg)):
-      raise TypeError("Argument '{}' of type {} is not a valid JAX type"
-                      .format(arg, type(arg)))
+def _check_arg(arg):
+  if not (isinstance(arg, core.Tracer) or _valid_jaxtype(arg)):
+    raise TypeError("Argument '{}' of type {} is not a valid JAX type"
+                    .format(arg, type(arg)))
 
 def _valid_jaxtype(arg):
   try:
diff --git a/jax/experimental/host_callback.py b/jax/experimental/host_callback.py
--- a/jax/experimental/host_callback.py
+++ b/jax/experimental/host_callback.py
@@ -193,14 +193,14 @@ def id_tap(func: Callable, arg, *,
   if func not in (_end_consumer, _unknown_testing_consumer):
     api._check_callable(func)
   flat_args, arg_treedef = pytree.flatten(arg)
-  api._check_args(flat_args)
+  for arg in flat_args: api._check_arg(arg)
   params = dict(kwargs)  #  we pass a copy of params to the primitive
   # See definition of id_tap_p for what parameters it takes
   params["func"] = func
   params["arg_treedef"] = arg_treedef
   if result is not None:
     flat_results, result_treedef = pytree.flatten(result)
-    api._check_args(flat_results)
+    for result in flat_results: api._check_arg(result)
     all_args = flat_args + flat_results
     params["nr_untapped"] = len(flat_results)
   else:


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -801,8 +801,49 @@ def test_grad_of_int_errors(self):
     dfn = grad(lambda x: x ** 2)
     self.assertRaisesRegex(
       TypeError,
-      "Primal inputs to reverse-mode differentiation must be of float or "
-      "complex type, got type int..", lambda: dfn(3))
+      (r"grad requires real- or complex-valued inputs \(input dtype that is a "
+       r"sub-dtype of np.floating or np.complexfloating\), but got int.*."),
+      lambda: dfn(3))
+
+  def test_grad_complex_result_errors(self):
+    dfn = grad(lambda x: x ** 2 + 1j)
+    self.assertRaisesRegex(
+      TypeError,
+      (r"grad requires real-valued outputs \(output dtype that is a "
+       r"sub-dtype of np.floating\), but got complex.*"),
+      lambda: dfn(3.))
+
+  def test_holomorphic_grad_of_float_errors(self):
+    dfn = grad(lambda x: x ** 2, holomorphic=True)
+    self.assertRaisesRegex(
+      TypeError,
+      (r"grad with holomorphic=True requires inputs with complex dtype, "
+       r"but got float.*"),
+      lambda: dfn(3.))
+
+  def test_holomorphic_jacrev_of_float_errors(self):
+    dfn = jacrev(lambda x: x ** 2, holomorphic=True)
+    self.assertRaisesRegex(
+      TypeError,
+      (r"jacrev with holomorphic=True requires inputs with complex dtype, "
+       r"but got float.*"),
+      lambda: dfn(3.))
+
+  def test_holomorphic_jacfwd_of_float_errors(self):
+    dfn = jacfwd(lambda x: x ** 2, holomorphic=True)
+    self.assertRaisesRegex(
+      TypeError,
+      (r"jacfwd with holomorphic=True requires inputs with complex dtype, "
+       r"but got float.*"),
+      lambda: dfn(3.))
+
+  def test_jacfwd_of_complex_errors(self):
+    dfn = jacfwd(lambda x: x ** 2)
+    self.assertRaisesRegex(
+      TypeError,
+      (r"jacfwd requires real-valued inputs \(input dtype that is a "
+       r"sub-dtype of np.floating\), but got complex.*"),
+      lambda: dfn(3. + 1j))
 
   def test_xla_computation(self):
     # these tests basically check the examples in the xla_computation docstring