jax.numpy.isin

目录

jax.numpy.isin#

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

确定 element 中的元素是否出现在 test_elements 中。

JAX 实现的 numpy.isin()

参数:
  • element (ArrayLike) – 输入要检查成员资格的元素数组。

  • test_elements (ArrayLike) – 用于检查每个元素是否存在的测试值的 N 维数组。

  • invert (bool) – 如果为真,返回 ~isin(element, test_elements)。默认为假。

  • assume_unique (bool) – 如果为真,则假设输入数组是唯一的,这可以导致更高效的计算。如果输入数组不是唯一的且 assume_unique 设置为 True,则结果是未定义的。

  • method – 指定用于计算结果的方法的字符串。支持的选项有 ‘compare_all’、’binary_search’、’sort’ 和 ‘auto’(默认)。

返回:

一个形状为 element.shape 的布尔数组,用于指定每个元素是否出现在 test_elements 中。

返回类型:

Array

示例

>>> 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)