Repository: google/jax
Issue: jnp.iscomplexobj calls __array__ on tracer objects

TLDR: In some exotic case `jnp.iscomplexobj(x)` calls `x.__array__()`, and if `x` happens to be a Tracer object then this errors out. This should never happen because the `dtype` of an array is a jit-time/constant information and so `array` should never be hit.

I noticed because in some complex NetKet code I attempted to add a check with  `jnp.iscomplexobj` to validate the input dtype. This works fine, but when this code is called within `jax.lax.custom_linear_solve` this crashes.

I suspect this is happening because  `jnp.iscomplexobj` is simply calling `np.iscomplexobj`, which probably has some fast-code path for some types (maybe those that define dtype), but if some condition is not met it calls `__array__` to convert it to a numpy array.

See the stack trace below, which you can also see [here](https://github.com/netket/netket/runs/4763942096?check_suite_focus=true#step:8:15062). I don't have a simple reproducer yet, unfortunately, but will attempt to derive one in the next days, unless some of you guys already have an idea of what is going wrong. 
The key point is when we hit some internal numpy frame

```python
test/optimizer/test_qgt_itersolve.py:119:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
netket/optimizer/linear_operator.py:100: in solve
    return self._solve(jax.tree_util.Partial(solve_fun), y, x0=x0, **kwargs)
netket/optimizer/qgt/qgt_jacobian_pytree.py:172: in _solve
    return _solve(self, solve_fun, y, x0=x0)
netket/optimizer/qgt/qgt_jacobian_pytree.py:251: in _solve
    out, info = solve_fun(unscaled_self, y, x0=x0)
../../../../../Documents/pythonenvs/netket_env/lib/python3.9/site-packages/jax/_src/scipy/sparse/linalg.py:686: in gmres
    x = lax.custom_linear_solve(A, b, solve=_solve, transpose_solve=_solve)
netket/optimizer/linear_operator.py:114: in __call__
    return self @ vec
netket/optimizer/qgt/qgt_jacobian_pytree.py:157: in __matmul__
    return _matmul(self, vec)
netket/optimizer/qgt/qgt_jacobian_pytree.py:203: in _matmul
    check_valid_vector_type(self.params, vec)
netket/optimizer/qgt/common.py:44: in check_valid_vector_type
    jax.tree_multimap(check, x, target)
netket/optimizer/qgt/common.py:32: in check                        #<---- x is a Tracer object. Maybe it has a weak dtype?
    if jnp.iscomplexobj(target) and not jnp.iscomplexobj(x):
<__array_function__ internals>:5: in iscomplexobj
    ???
../../../../../Documents/pythonenvs/netket_env/lib/python3.9/site-packages/numpy/lib/type_check.py:316: in iscomplexobj
    type_ = asarray(x).dtype.type
../../../../../Documents/pythonenvs/netket_env/lib/python3.9/site-packages/numpy/core/_asarray.py:102: in asarray
    return array(a, dtype, copy=False, order=order)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = TracerArrayConversionError(Traced<ShapedArray(float64[5])>with<DynamicJaxprTrace(level=1/2)>), tracer = Traced<ShapedArray(float64[5])>with<DynamicJaxprTrace(level=1/2)>

    def __init__(self, tracer: "core.Tracer"):
      super().__init__(
          "The numpy.ndarray conversion method __array__() was called on "
>         f"the JAX Tracer object {tracer}{tracer._origin_msg()}")
E     IndexError: list assignment index out of range
```

The dtypes of `x` and `target` on that highlighted line are
```python
x=Traced<ShapedArray(complex128[5])>with<DynamicJaxprTrace(level=1/2)>  
target=(Traced<ShapedArray(float64[5])>with<DynamicJaxprTrace(level=1/2)>, Traced<ShapedArray(float64[5])>with<DynamicJaxprTrace(level=1/2)>)
```

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 84a0565f9..6d80ec3d4 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -132,7 +132,13 @@ get_printoptions = np.get_printoptions
 printoptions = np.printoptions
 set_printoptions = np.set_printoptions
 
-iscomplexobj = np.iscomplexobj
+@_wraps(np.iscomplexobj)
+def iscomplexobj(x):
+  try:
+    typ = x.dtype.type
+  except AttributeError:
+    typ = asarray(x).dtype.type
+  return issubdtype(typ, complexfloating)
 
 shape = _shape = np.shape
 ndim = _ndim = np.ndim
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index e3477d635..0e2691a76 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4134,6 +4134,14 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
         np.array(bytearray(b'\\x2a\\xf3'), ndmin=2)
     )
 
+  @parameterized.named_parameters(
+    {'testcase_name': f"_val={val}", 'val': val}
+    for val in [1+1j, [1+1j], jnp.pi, np.arange(2)])
+  def testIsComplexObj(self, val):
+    args_maker = lambda: [val]
+    self._CheckAgainstNumpy(np.iscomplexobj, jnp.iscomplexobj, args_maker)
+    self._CompileAndCheck(jnp.iscomplexobj, args_maker)
+
   def testIsClose(self):
     c_isclose = jax.jit(jnp.isclose)
     c_isclose_nan = jax.jit(partial(jnp.isclose, equal_nan=True))