jax.lax.select

jax.lax.select(pred, on_true, on_false)

Wraps XLA’s Select operator.

Parameters

• pred (Union[Any, ndarray, bool_, number, bool, int, float, complex]) –

• on_true (Union[Any, ndarray, bool_, number, bool, int, float, complex]) –

• on_false (Union[Any, ndarray, bool_, number, bool, int, float, complex]) –

Return type

Any