AttributeError: 'Unit' object has no attribute 'shape' while tracing function. What does it mean??

[RadostW]

Inside a nested scan when using jax.lax.cond I get error message that is hard to understand.
Original code was very complex, sadly minimal example I could make is still quite complex.
Bizzarly, when changing small details code runs fine (see commented lines).

Working fine:
x * d_www[0,0,0] # Ok
x * f_www(x)[0,0,0,0] # Ok

Throwing hard to understand error
jnp.array([1.0]) * f_www(x)[0, 0, 0, 0] * d_www[0, 0, 0] # AttributeError

Which makes me think that both d_www creation and f_www creation work fine.

I've found no similar issues when searching.

~~R

import jax.numpy as jnp
import jax


def upper(y):
    return jnp.exp(y)


def lower(y):
    return jnp.exp(2 * y)


def branch_fun(q):
    x = q[0]
    ret = jax.lax.cond(x > 1.0, upper, lower, x)
    return jnp.array([[ret]])


def L_w(f):
    def wrapped(x):
        return jax.numpy.tensordot(jax.jacobian(f)(x), branch_fun(x), axes=1)

    return wrapped


def solve_many():
    t0 = 0.0
    w0 = jax.numpy.zeros(1)

    id_ = lambda x: x
    f_www = L_w(L_w(L_w(id_)))

    def step(
        x,
        d_www,
    ):

        new_x = (
            # x * d_www[0,0,0] # Ok
            # x * f_www(x)[0,0,0,0] # Ok
            jnp.array([1.0]) * f_www(x)[0, 0, 0, 0] * d_www[0, 0, 0]  # AttributeError
            # x * f_www(x)[0,0,0,0] * d_www[0,0,0] # AttributeError
            # contract_all(f_www(x), d_www) # AttributeError
            # f_www(x)[:,0,0,0] # AttributeError
        )

        return new_x

    key = jax.random.PRNGKey(0)

    def scan_func(x, y):
        xp = step(x, y)
        return (xp, xp)

    def chunk_function(x, y):
        z = jax.lax.scan(scan_func, jnp.array([0.1]), y)[0]
        return z, z

    def get_solution_fragment(starting_state, key):
        wiener_integrals = jax.random.normal(key, shape=(4, 1, 1, 1))

        last_state, solution_values = jax.lax.scan(
            chunk_function,
            starting_state,
            jnp.reshape(wiener_integrals, (2, 2, 1, 1, 1)),
        )

        return (last_state, last_state)

    @jax.vmap
    def get_solution(key):
        _, chunked_solution = jax.lax.scan(
            lambda state, key: get_solution_fragment(state, key),
            (jnp.array([0.1])),
            jax.random.split(key, 1),
        )

        return chunked_solution

    keys = jax.random.split(key, 1)
    solutions = get_solution(keys)

    return solutions


def test_branched_coefficients():
    trajectories = solve_many()


if __name__ == "__main__":
    trajectories = solve_many()

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_branched_coefficients.py", line 91, in <module>
    trajectories = solve_many()
  File "test_branched_coefficients.py", line 81, in solve_many
    solutions = get_solution(keys)
  File "test_branched_coefficients.py", line 72, in get_solution
    _, chunked_solution = jax.lax.scan(
AttributeError: 'Unit' object has no attribute 'shape'

[YouJiacheng]

I get a minimal repro:

import jax.numpy as jnp
import jax

id_f = lambda x: x
def b(x):
    return jax.lax.cond(x[0], id_f, id_f, x) # jnp.where(x[0], id_f(x), id_f(x)) is ok

f = jax.jacrev(lambda x: b(b(x))) # jacfwd is ok

def repro():
    jax.vmap(f)(jnp.ones((1,1)))

def ok():
    f(jnp.ones((1,)))
    print('ok')

if __name__ == "__main__":
    with jax.disable_jit():
        ok()
        repro()

[RadostW]

The jacrev -> jacfwd and hessian -> jacfwd(jacfwd(f)) fixed my original issue (at a big cost to performance), thanks! 😄
I would still rather use jacrev in some places for speed though.

[YouJiacheng]

If you need a workaround, using jnp.where or lax.select to replace lax.cond may have a much smaller performance overhead!

[mattjj]

Thanks, this is definitely a bug! Actually it might fit better as a bug issue than a 'discussion'.

[mattjj]

Let's follow up on #9985.