equivalent of jnp.where in jit functions

[agobL]

Hello,

I want to use the following code in a function using jit :

indexes = np.where(x < 0)
y = np.ones(8135)
y[indexes] = 0

However jnp.where provides an error with abstract values. So I used the jax.lax.cond function but with that my code is slower than before. I would like to know if there is another way to programm these 3 lines without using jax.lax.cond.

Thanks in advance.

[jakevdp]

In many cases, you an use the jit-compatible three-term jnp.where instead. Here it would look like:

y = jnp.where(x < 0, 0, 1)

There are some more examples & discussion in the documentation associated with the error that results from your original code: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

[agobL]

Thank you for the answer. It solves my issue.

[nt54hamnghi]

Hi @jakevdp,
I thought if x is a positional argument, it would be traced, thus couldn't be used in a control-flow statement. How is it possible to use x in jnp.where? Does it have something to do with the underlying implementation of jnp.where?

[jakevdp]

That statement is true about control flow like if..else statements. jnp.where is different: it is fine to pass traced values so long as you're using the three argument version mentioned above.

[nt54hamnghi]

Thanks for your reply. It makes sense to me now.