带有任意运算符的 numpy 点或 einsum

Question

我想使用类似 np.dot 或(最好)np.einsum 的东西来有效地执行它们相同的功能，但使用备用 ufunc 而不是 np.multiply。例如，考虑这两个数组:

>>> a
array([[0, 1],
       [1, 1],
       [1, 0]])
>>> b
array([[0, 0],
       [1, 0],
       [1, 0],
       [0, 0]])

现在假设我想计算 a 每一行中元素的数量等于 b 每一行中相应元素的数量。我希望能够执行以下等效操作(注意:下面的输出是捏造的，但值是我希望看到的值):

>>> np.dot(a, b.T, ufunc=np.equal)
array([[1, 0, 0, 1],
       [0, 1, 1, 0],
       [1, 2, 2, 1]])

有办法吗？

Answer 1

您可以将 Divakar 的回答中的广播与 numexpr 一起使用:

numexpr.evaluate('sum(1*(a == b), axis=2)', {'a': a[:,None]})

1*() 是 a workaround.我已经确认这不会分配一个大的临时数组。