jax.numpy.isin#

jax.numpy.isin(element,test_elements,assume_unique=False,invert=False,*,method='auto')[source]#

Determine whether elements inelement appear intest_elements.

JAX implementation ofnumpy.isin().

Parameters:

• element (ArrayLike) – input array of elements for which membership will be checked.

• test_elements (ArrayLike) – N-dimensional array of test values to check for the presence ofeach element.

• invert (bool) – If True, return~isin(element,test_elements). Default is False.

• assume_unique (bool) – if true, input arrays are assumed to be unique, which canlead to more efficient computation. If the input arrays are not uniqueand assume_unique is set to True, the results are undefined.

• method – string specifying the method used to compute the result. Supportedoptions are ‘compare_all’, ‘binary_search’, ‘sort’, and ‘auto’ (default).

Returns:

A boolean array of shapeelement.shape that specifies whether each elementappears intest_elements.

Return type:

Array

Examples

>>>elements=jnp.array([1,2,3,4])>>>test_elements=jnp.array([[1,5,6,3,7,1]])>>>jnp.isin(elements,test_elements)Array([ True, False,  True, False], dtype=bool)