二维数组的按行 numpy.isin [重复]

我有两个数组：

A = np.array([[3, 1], [4, 1], [1, 4]])
B = np.array([[0, 1, 5], [2, 4, 5], [2, 3, 5]])

是否可以使用numpy.isin二维数组按行排列？我想检查一下是否A[i,j] is in B[i]并将结果返回到C[i,j]。最后我会得到以下内容C:

np.array([[False, True], [True, False], [False, False]])

如果这也可以通过==运算符，那么我也可以将它与 PyTorch 一起使用。

Edit:我也考虑过检查不同 numpy 数组中的相同行 https://stackoverflow.com/q/51352527/11227118。这个问题在某种程度上是相似的，但我无法将其解决方案应用于这个略有不同的问题。

不确定我的代码是否完美解决您的问题。请在更多测试用例上运行它来确认。但我会像我在下面的代码中所做的那样，利用 numpys 向量外运算能力（类似于向量外积）。如果它按预期工作，它也应该与 pytorch 一起工作。

import numpy as np
A = np.array([[3, 1], [4, 1], [1, 4]])
B = np.array([[0, 1, 5], [2, 4, 5], [2, 3, 5]])

AA = A.reshape(3, 2, 1)
BB = B.reshape(3, 1, 3)
(AA == BB).sum(-1).astype(bool)

output:

array([[False,  True],
       [ True, False],
       [False, False]])