lax.select same shape constraint

[LeonEricsson]

I switched over some of my jnp.where calls to jax.lax.select and noticed it was more restrictive in terms of shapes. In the code snippet below, timer is a Traced<ShapedArray(int32[1]).

x = jax.lax.select(
            timer > 0,
            True,
            False,
        )

this will fail because True and False are scalars. Is the below code the most straightforward way of solving this? It seems excessive?

x = jax.lax.select(
            x > 0,
            jnp.ones((1,), dtype=jnp.bool),
            jnp.zeros((1,), dtype=jnp.bool),
        )

is this case simply more suited for jax.numpy.where given that it will broadcast?

EDIT: Another question, how does one deal with the scenario where your condition is a simple scalar and the cases are arrays of the same shape?

[jakevdp]

You'll find that functions in jax.lax are generally more restrictive than equivalent functions in jax.numpy, particularly when it comes to implicit broadcasting, type promotion, and rank promotion. If you want to use lax functions, you'll generally have to do manual dtype and shape conversions to ensure the inputs match the more restrictive APIs. If you don't want the more restrictive API of lax, then you should use jax.numpy.

Another question, how does one deal with the scenario where your condition is a simple scalar and the cases are arrays of the same shape?

It sounds like you might be looking for lax.cond, which is only valid for scalar conditions:

from jax import lax
x = 1
lax.cond(x > 0, lambda: True, lambda: False)
# Array(True, dtype=bool)

[LeonEricsson]

'preciate it thanks!