jax.numpy.isin

内容

jax.numpy.isin#

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

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

JAX 对 numpy.isin() 的实现。

参数：:

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

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

  • invert (bool) – 如果为 True，则返回 ~isin(element, test_elements)。默认为 False。

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

  • method – 指定用于计算结果的方法的字符串。支持的选项为 'compare_all'、'binary_search'、'sort' 和 'auto'（默认）。

返回:

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

返回类型:

数组

示例

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