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)>)
```


Discussion:
When was this fixed? I'm still having this issue.
I'm not aware of any fix related to this bug, because it was closed shortly after being reported. Looking into it now, it seems that this may come up in cases where collections like tuples or lists are passed to `jnp.iscomplexobj` within a `jit` context. For example:
```python
import jax

@jax.jit
def f(x):
  return jnp.iscomplexobj(x)

f([1])
```
```pytb
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/numpy/lib/type_check.py](https://localhost:8080/#) in iscomplexobj(x)
    336     try:
--> 337         dtype = x.dtype
    338         type_ = dtype.type

AttributeError: 'list' object has no attribute 'dtype'

During handling of the above exception, another exception occurred:

UnfilteredStackTrace                      Traceback (most recent call last)
20 frames
<__array_function__ internals> in iscomplexobj(*args, **kwargs)

UnfilteredStackTrace: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
While tracing the function f at <ipython-input-17-86c68a464777>:3 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TracerArrayConversionError                Traceback (most recent call last)
<__array_function__ internals> in iscomplexobj(*args, **kwargs)

[/usr/local/lib/python3.7/dist-packages/numpy/lib/type_check.py](https://localhost:8080/#) in iscomplexobj(x)
    338         type_ = dtype.type
    339     except AttributeError:
--> 340         type_ = asarray(x).dtype.type
    341     return issubclass(type_, _nx.complexfloating)
    342 

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
While tracing the function f at <ipython-input-17-86c68a464777>:3 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
```
Does that seem similar to the issue you're having?
Reopening, because I think it's worth fixing `jnp.iscomplexobj` to handle this case.
I'll try to create a MWE to recreate the error I'm seeing, I noticed it when trying to add support for complex numbers to diffrax - https://github.com/patrick-kidger/diffrax/pull/112
Looking back at this – it looks like similar issues arise for other `jax.numpy` functions that are aliased to numpy equivalents. For example:
```python
import jax.numpy as jnp
from jax import jit

jit(jnp.iscomplexobj)([1, 2, 3])  # TracerArrayConversionError
jit(jnp.shape)([1, 2, 3])  # TracerArrayConversionError
jit(jnp.size)([1, 2, 3])  # TracerArrayConversionError
jit(jnp.ndim)([1, 2, 3])  # TracerArrayConversionError
```
We should probably make these match other JAX APIs and wrap these in such a way that `_check_arraylike` is called on the inputs.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -132,7 +132,13 @@ def canonicalize_shape(
 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


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4134,6 +4134,14 @@ def testMemoryView(self):
         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))