多维jax.isin()

Question

我正在尝试过滤一个三元组的数组。我要筛选的标准是，另一个三元组是否至少包含一个具有相同的第一个和第三个元素的元素。E.g

import jax.numpy as jnp
array1 = jnp.array(
  [
    [0,1,2],
    [1,0,2],
    [0,3,3],
    [3,0,1],
    [0,1,1],
    [1,0,3],
  ]
)
array2 = jnp.array([[0,1,3],[0,3,2]])
# the mask to filter the first array1 should look like this:
jnp.array([True,False,True,False,False,False])

使用jax实现此掩码的有效计算方法是什么？我期待着你的意见。

Answer 1

要做到这一点，可以减少通过广播的平等检查：

import jax.numpy as jnp
array1 = jnp.array(
  [
    [0,1,2],
    [1,0,2],
    [0,3,3],
    [3,0,1],
    [0,1,1],
    [1,0,3],
  ]
)
array2 = jnp.array([[0,1,2],[0,3,2]])  # note adjustment to match first entry of array1

mask = (array1[:, None] == array2[None, :]).all(-1).any(-1)
print(mask)
# [ True False False False False False]

XLA没有任何类似于二进制搜索的原语，所以通常最好的方法是生成完全相等的矩阵并进行约简。如果您像GPU/TPU一样在加速器上运行代码，这种向量化操作是高效的并行化操作，因此在实践中将非常有效地计算它。