using traced array in lax.cond()

[LindoNkambule]

Hi,

I have this function

@jit
def my_func(x, alpha, lam):
    return -alpha*(jnp.cosh(x)-1.0)-lam*(jnp.exp(x)-x-1.0)

which I call in another function

x = -my_func(1.0, alpha, lam)

is it possible to use x in lax.cond as follows (or is there a workaround):

t = lax.cond(
        (x <0.5),
        lambda: jnp.float32(1),
        lambda: jnp.sqrt(jnp.float32(2)/(alpha+lam))
)

?

Thank you

[mattjj]

Thanks for the question!

If I understand the question correctly, yes that should work. (But also based only on what you said a simple t = jnp.float32(1) if x < 0.5 else jnp.sqrt(jnp.float32(2)/(alpha+lam)) would work too, unless everything is under an outer jit, in which case you'd need a lax.cond again. That's basically what lax.cond is there for: it always works regardless of what tracing/transformations are going on.)

[LindoNkambule]

Thank you!