TypeError: unhashable type: 'DeviceArray'

[dui1234]

Hi,

I am trying to code a body of a while loop with some conditions in its body function as folloiwng:

def body_fun(x,y):
    left = 2 * x
    right = left + 1
    cond = x > y
    idx = jnp.where(cond, left, right)
    return idx
    
init_val = 0
y = 10
cond_fun = lambda x: x<10
lax.while_loop(cond_fun, body_fun(init_val,y), init_val)

However, this returns TypeError: unhashable type: 'DeviceArray'. How should I fix this?

Thanks,
Aom

[jakevdp]

The issue is you need to pass body_fun to your while loop, rather than the result of evaluating body_fun. Additionally, both body_fun and cond_fun must accept a single value as an argument. It sounds like what you want is to partially evaluate body_fun before passing it to while_loop, e.g.

from functools import partial
lax.while_loop(cond_fun, partial(body_fun, y=y), init_val)

[dui1234]

Very much obliged! Your explaination is very helpful!

[jakevdp]

We're improving the errors for these cases in #10262