Repository: google/jax

Issue:
`float()` argument `TypeError` in a scan.
### Discussed in https://github.com/google/jax/discussions/19057

<div type='discussions-op-text'>

<sup>Originally posted by **xaviergonzalez** December 19, 2023</sup>
I have a really weird jax bug with scan, I'm wondering if anyone else has seen similar issues?
Here's a minimal reproducible example:
```
import jax
import jax.numpy as jnp
import jax.random as jr

def g(hiddens):
    hiddens_aug = jnp.vstack((hiddens[0], hiddens))
    new_hiddens = hiddens_aug.copy()
    diff = new_hiddens[:-1] - hiddens
    diff = new_hiddens[:-1] - hiddens
    out = jnp.trace(jnp.conj(diff).T @ diff).real
    return jnp.array(out, dtype=jnp.complex64)


def _step(carry, arg):
    primals, f_vjp = jax.vjp(
        g,
        jr.normal(jr.PRNGKey(0), (9, 8), dtype=jnp.complex64),
    )
    out = f_vjp(1.0 + 0j)[0]
    return carry, carry
```
So far, so good. I've defined a function g, that basically maps from \C^8 to \R (casted to be complex for typing reasons), and used vjp to differentiate it (grad behaves a bit weird with these functions). Now, I'd just like to do a lot of these iterations with a scan.
The following code works completely fine
```
primals, f_vjp = jax.vjp(g, jr.normal(jr.PRNGKey(0), (9,8), dtype=jnp.complex64))
out = f_vjp(1.0 + 0j)[0]
```
out has shape (9,8), and dtype complex64, as desired.

Note that this code are basically the two lines of the _step, except they aren't even returned.

However, the following code returns a TypeError
```
a, b = jax.lax.scan(_step, 0, jnp.arange(4, dtype=jnp.complex64))
```

Gives error:

```
TypeError                                 Traceback (most recent call last)
<ipython-input-4-e5de119df33f> in <cell line: 1>()
----> 1 a, b = jax.lax.scan(_step, 0, jnp.arange(4, dtype=jnp.complex64))

    [... skipping hidden 10 frame]

/usr/local/lib/python3.10/dist-packages/jax/_src/lax/lax.py in _convert_elt_type_folding_rule(consts, eqn)
   2375       isinstance(o.aval, core.UnshapedArray) and not np.shape(c) and
   2376       not dtypes.issubdtype(eqn.params['new_dtype'], dtypes.extended)):
-> 2377     out = np.array(c, eqn.params['new_dtype'])
   2378     if not o.aval.weak_type:
   2379       return [out], None
TypeError: float() argument must be a string or a real number, not 'complex'
```

Notice that `_step` doesn't even return anything relating to its intermediate computations! Moreover, I can see that when I comment out the line
```
out = f_vjp(1.0 + 0j)[0]
```
the error goes away.

I'm really at a loss as to why this line is completely fine outside of a scan, but crashes the scan, even though it is unrelated to either the input or output of `_step`. I've really been struggling to debug inside of the scan.
If anyone has seen this type of error before, or even just has tips about how to debug inside of a scan, I'd be super grateful to hear them.</div>


Discussion:
We're basically doing this inside a constant folding step:
```python
import numpy as np
np.array(1+0j, np.dtype('float32'))
```

which gives

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[2], line 1
----> 1 np.array(1+0j, np.dtype('float32'))

TypeError: float() argument must be a string or a real number, not 'complex'
```

So our constant folding logic is wrong, just in that apparently NumPy doesn't like when we try to ask for a float32 array containing a complex (builtin) type.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2382,7 +2382,9 @@ def _convert_elt_type_folding_rule(consts, eqn):
   if (type(c) in {np.ndarray, *dtypes.python_scalar_dtypes} and
       isinstance(o.aval, core.UnshapedArray) and not np.shape(c) and
       not dtypes.issubdtype(eqn.params['new_dtype'], dtypes.extended)):
-    out = np.array(c, eqn.params['new_dtype'])
+    with warnings.catch_warnings():
+      warnings.simplefilter('ignore', np.ComplexWarning)
+      out = np.array(c).astype(eqn.params['new_dtype'])
     if not o.aval.weak_type:
       return [out], None
     out = out.item()


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2712,6 +2712,27 @@ def f(x):
 
     jax.hessian(f)(1.0)  # don't crash
 
+  def test_constant_folding_complex_to_real_scan_regression(self):
+    # regression test for github.com/google/jax/issues/19059
+    def g(hiddens):
+        hiddens_aug = jnp.vstack((hiddens[0], hiddens))
+        new_hiddens = hiddens_aug.copy()
+        diff = new_hiddens[:-1] - hiddens
+        diff = new_hiddens[:-1] - hiddens
+        out = jnp.trace(jnp.conj(diff).T @ diff).real
+        return jnp.array(out, dtype=jnp.complex64)
+
+
+    def _step(carry, arg):
+        primals, f_vjp = jax.vjp(
+            g,
+            jax.random.normal(jax.random.key(0), (9, 8), dtype=jnp.complex64),
+        )
+        out = f_vjp(np.array(1.0 + 0j, 'complex64'))[0]
+        return carry, carry
+
+    a, b = jax.lax.scan(_step, 0, jnp.arange(4, dtype=jnp.complex64))
+
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):