numpy dot or einsum with arbitrary operator

Question

I would like to use something like np.dot or (preferably) np.einsum to efficiently perform their same function but with an alternate ufunc instead of np.multiply. For example, consider these two arrays:

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

Now suppose I want to count the number of elements in each row of a equal to the corresponding elements in each row of b. I'd like to be able to do the equivalent of the following (note: the output below is fabricated but the values are what I would expect to see):

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

Is there a way to do this?

Answer 1

You can use the broadcasting from Divakar's answer together with numexpr:

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

The 1*() is a workaround. I have confirmed this doesn't allocate a big temporary array.

Answer 2

You could use broadcasting for such match counting problem -

(a[:,None] == b).sum(2)

Sample run -

In [36]: a
Out[36]: 
array([[0, 1],
       [1, 1],
       [1, 0]])

In [37]: b
Out[37]: 
array([[0, 0],
       [1, 0],
       [1, 0],
       [0, 0]])

In [38]: (a[:,None] == b).sum(2)
Out[38]: 
array([[1, 0, 0, 1],
       [0, 1, 1, 0],
       [1, 2, 2, 1]])

If you really want to employ np.einsum and np.equal, here's a way to mold the earlier approach to give us the desired result -

np.einsum('ijk->ij',np.equal(a[:,None],b).astype(int))