Multiple conditions in control flow

[ToshiyukiBandai]

Hi all,

I would like to implement a jitted indicator function with multiple conditions. If I put two conditions in jnp.where, it gives an error (TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..):

import jax.numpy as jnp
from jax jit

@jit
def f(x):
    print("compile")
    return jnp.where(x < 3 and x > 1, 1.0, 0.0)

I was able to get around this by separating the two conditions:

@jit
def f(x):
    print("compile")
    I = jnp.where(x < 3, 1.0, 0.0)
    I = jnp.where(x > 1, 1.0, 0.0)
    return I

Could anyone explain what is going on here? Is there any other ways better than my second implementation?

[jakevdp]

Thanks for the question! Instead of x < 3 and x > 1, you should use (x < 3) & (x > 1). The reason this is required is that there is no way to overload the and operator in Python, so we follow NumPy's convention of using bitwise & for element-wise boolean logic.

Note that due to Python's operator precedence conventions, the parentheses in the expression above cannot be omitted.

Similarly, in this context you would use | instead of or, ~ instead of not, and as a bonus you get ^ for XOR, which has no python keyword equivalent.

[ToshiyukiBandai]

Oh, I see. I tried (x < 3) and (x > 1) but did not notice and was the issue. Thank you so much for the clarification!

[patrick-kidger]

Although by the way, I would recommend against using ~. If you happen to have a raw Python True or False rather than a jnp.array(...), then you will get different behaviour:

~True  # -2
~False  # -1

so for safety's sake it is best to always use jnp.invert instead.