float() argument TypeError in a scan.

[xaviergonzalez]

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.

[mattjj]

Sounds like an internal JAX bug! On that presumption, let me convert this to a bug issue...

Let's discuss on #19059 (rather than here).

[xaviergonzalez]

FWIW, the error goes away if I just return out instead of trying to cast it to a complex number. I.e., the following code is fine:

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 out


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
    
a, b = jax.lax.scan(_step, 0, jnp.arange(4))

I don't totally understand why this makes a difference, though.