jax.numpy.select

jax.numpy.select(condlist, choicelist, default=0)

Return an array drawn from elements in choicelist, depending on conditions.

LAX-backend implementation of numpy.select().

Original docstring below.

Parameters

• condlist (list of bool ndarrays) – The list of conditions which determine from which array in choicelist the output elements are taken. When multiple conditions are satisfied, the first one encountered in condlist is used.

• choicelist (list of ndarrays) – The list of arrays from which the output elements are taken. It has to be of the same length as condlist.

• default (scalar, optional) – The element inserted in output when all conditions evaluate to False.

Returns

output – The output at position m is the m-th element of the array in choicelist where the m-th element of the corresponding array in condlist is True.

Return type

ndarray