Repository: google/jax

Issue:
An error from a combination of scan, complex numbers, and np.diag
A very high level summary: I got an error when I did some complex number computation inside a lax.scan function (compute1 in the code below). But the almost identical code can be run with a for loop (compute2). Only gradient computation raises the error. Forward computation is fine for both method. If I replace np.diag(d) with v in the function f, then also no error.

```python
import jax.numpy as np
from jax import grad, random, jit, lax

def f(carrier, x):
    h =  x*np.array([[0, 1j], [-1j, 0]])
    d, v = np.linalg.eigh(h, symmetrize_input=False)
    carrier = np.diag(d) @ carrier
    return carrier, None

def compute1(p):
    u = np.eye(2, dtype=np.complex64)
    carrier, _ = lax.scan(f, u, p)
    return np.abs(np.trace(carrier))

def compute2(p):
    u = np.eye(2, dtype=np.complex64)
    for i in range(3):
        h =  p[i]*np.array([[0, 1j], [-1j, 0]])

        d, v = np.linalg.eigh(h, symmetrize_input=False)
        u = np.diag(d) @ u
    return np.abs(np.trace(u))

print(grad(compute1)(np.arange(2,5,dtype=np.complex64)))
# print(compute1(np.arange(2,5,dtype=np.complex64)))
```

Error message:
/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
  File "/home/xiaotongni/jax-qoc/test_scan.py", line 24, in <module>
    print(grad(compute1)(np.arange(2,5,dtype=np.complex64)))
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/api.py", line 352, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/api.py", line 407, in value_and_grad_f
    ans, vjp_py = vjp(f_partial, *dyn_args)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/api.py", line 1285, in vjp
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 106, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 95, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/linear_util.py", line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/xiaotongni/jax-qoc/test_scan.py", line 12, in compute1
    carrier, _ = lax.scan(f, u, p)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lax/lax_control_flow.py", line 804, in scan
    linear=(False,) * (len(consts) + len(in_flat)))
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lax/lax_control_flow.py", line 1166, in scan_bind
    num_carry=num_carry, linear=linear)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/core.py", line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 305, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lax/lax_control_flow.py", line 865, in _scan_jvp
    jaxpr, nonzeros, instantiate=carry_nz + [False] * num_ys)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 575, in jvp_jaxpr
    jaxpr_out, pvals_out, literals_out = pe.trace_to_jaxpr(f_jvp, pvals, instantiate=True)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/linear_util.py", line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/core.py", line 97, in jaxpr_as_fun
    return eval_jaxpr(typed_jaxpr.jaxpr, typed_jaxpr.literals, *args)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/core.py", line 251, in eval_jaxpr
    ans = eqn.primitive.bind(*(subfuns + in_vals), **params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/core.py", line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 305, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 388, in linear_jvp
    val_out = primitive.bind(*primals, **params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/core.py", line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lax/lax.py", line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
  File "/home/xiaotongni/miniconda3/lib/python3.7/site-packages/jax/lax/lax.py", line 2484, in _pad_dtype_rule
    raise TypeError(msg.format(operand.dtype, padding_value.dtype))
TypeError: pad operand and padding_value must be same dtype: got complex64 and float32.


Discussion:

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -265,10 +265,10 @@ def eigh_jvp_rule(primals, tangents, lower):
   a, = primals
   a_dot, = tangents
 
-  v, w = eigh_p.bind(symmetrize(a), lower=lower)
+  v, w_real = eigh_p.bind(symmetrize(a), lower=lower)
 
   # for complex numbers we need eigenvalues to be full dtype of v, a:
-  w = w.astype(a.dtype)
+  w = w_real.astype(a.dtype)
   eye_n = jnp.eye(a.shape[-1], dtype=a.dtype)
   # carefully build reciprocal delta-eigenvalue matrix, avoiding NaNs.
   Fmat = jnp.reciprocal(eye_n + w[..., jnp.newaxis, :] - w[..., jnp.newaxis]) - eye_n
@@ -277,8 +277,8 @@ def eigh_jvp_rule(primals, tangents, lower):
                 precision=lax.Precision.HIGHEST)
   vdag_adot_v = dot(dot(_H(v), a_dot), v)
   dv = dot(v, jnp.multiply(Fmat, vdag_adot_v))
-  dw = jnp.diagonal(vdag_adot_v, axis1=-2, axis2=-1)
-  return (v, w), (dv, dw)
+  dw = jnp.real(jnp.diagonal(vdag_adot_v, axis1=-2, axis2=-1))
+  return (v, w_real), (dv, dw)
 
 def eigh_batching_rule(batched_args, batch_dims, lower):
   x, = batched_args


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -154,6 +154,15 @@ def check_close(xs, ys, atol=None, rtol=None):
   assert_close = partial(_assert_numpy_close, atol=atol, rtol=rtol)
   tree_all(tree_multimap(assert_close, xs, ys))
 
+def _check_dtypes_match(xs, ys):
+  def _assert_dtypes_match(x, y):
+    if FLAGS.jax_enable_x64:
+      assert _dtype(x) == _dtype(y)
+    else:
+      assert (dtypes.canonicalize_dtype(_dtype(x)) ==
+              dtypes.canonicalize_dtype(_dtype(y)))
+  tree_all(tree_multimap(_assert_dtypes_match, xs, ys))
+
 
 def inner_prod(xs, ys):
   def contract(x, y):
@@ -202,7 +211,9 @@ def check_jvp(f, f_jvp, args, atol=None, rtol=None, eps=EPS):
   rng = np.random.RandomState(0)
   tangent = tree_map(partial(rand_like, rng), args)
   v_out, t_out = f_jvp(args, tangent)
+  _check_dtypes_match(v_out, t_out)
   v_out_expected = f(*args)
+  _check_dtypes_match(v_out, v_out_expected)
   t_out_expected = numerical_jvp(f, args, tangent, eps=eps)
   # In principle we should expect exact equality of v_out and v_out_expected,
   # but due to nondeterminism especially on GPU (e.g., due to convolution
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -418,7 +418,7 @@ def testStaticIndexingGrads(self, shape, dtype, rng_factory, indexer):
     rng = rng_factory(self.rng())
     tol = 1e-2 if jnp.finfo(dtype).bits == 32 else None
     arg = rng(shape, dtype)
-    fun = lambda x: x[indexer]**2
+    fun = lambda x: jnp.asarray(x)[indexer]**2
     check_grads(fun, (arg,), 2, tol, tol, tol)
 
   def _ReplaceSlicesWithTuples(self, idx):
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -405,6 +405,8 @@ def testEighGradVectorComplex(self, shape, dtype, rng_factory, lower, eps):
     # evaluate eigenvector gradient and groundtruth eigensystem for perturbed input matrix
     f = partial(jnp.linalg.eigh, UPLO=uplo)
     (w, v), (dw, dv) = jvp(f, primals=(a,), tangents=(a_dot,))
+    self.assertTrue(jnp.issubdtype(w.dtype, jnp.floating))
+    self.assertTrue(jnp.issubdtype(dw.dtype, jnp.floating))
     new_a = a + a_dot
     new_w, new_v = f(new_a)
     new_a = (new_a + np.conj(new_a.T)) / 2