CODEWITHSUNDEEP
pythonnumpynumpy-einsum
Miko Miko
Miko Miko

Reputation: 11

Standard operations equivalent of einsum expression

I have the following einsum expressions:

np.einsum("abc,ab->ac",a,b)

np.einsum("abc,abd->dc", a, b)

That I would need to convert to standard numpy operations. Can anyone help me?

Upvotes: 1

Views: 119

Answers (1)

hpaulj
hpaulj

Reputation: 231550

np.einsum("abc,ab->ac",a,b) treats a as the batch dimension, carried through unchanged. b is the sum-of-products dimension.

(a*b[:,:,None]).sum(axis=1)

should be one equivalent

(a.transpose(0,2,1) @ b[:,:,None])[:,:,0]

should also work - but I'll need to run a test case to be sure.

np.einsum("abc,abd->dc", a, b) does sum-of-products on 'ab', with 'dc' reversed

(a[:,:,None,:]*b[:,:,:,None]).sum(axis=(0,1))

forms a 'abdc' product, and then sums.

 b.reshape(-1,b.shape[-1]).transpose() @ a.reshape(-1,a.shape[-1])

is my guess at the matmul version. reshape reduces the 'ab' dimensions to one.

The matmul equivalents have to have the form

(B,N,M) @ (B,M,O) -> (B,N,O)

with B as batch, M as sum-of-products

test

In [7]: a,b,c,d = 2,3,4,5
In [8]: A = np.arange(a*b*c).reshape(a,b,c); B = np.arange(a*b).reshape(a,b)
In [9]: C = np.einsum('abc,ab->ac',A,B)
In [10]: C.shape
Out[10]: (2, 4)

In [11]: D = (A*B[:,:,None]).sum(axis=1)
In [12]: D.shape
Out[12]: (2, 4)

In [13]: np.allclose(C,D)
Out[13]: True

In [16]: E=(A.transpose(0,2,1) @ B[:,:,None])[:,:,0]
In [17]: E.shape
Out[17]: (2, 4)
In [18]: np.allclose(C,E)
Out[18]: True

and

In [19]: B=np.arange(a*b*d).reshape(a,b,d)
In [20]: C=np.einsum("abc,abd->dc", A,B)
In [21]: C.shape
Out[21]: (5, 4)

In [22]: D=(A[:,:,None,:]*B[:,:,:,None]).sum(axis=(0,1))
In [23]: D.shape
Out[23]: (5, 4)
In [24]: np.allclose(C,D)
Out[24]: True

In [25]: E=B.reshape(-1,B.shape[-1]).transpose() @ A.reshape(-1,A.shape[-1])
In [26]: E.shape
Out[26]: (5, 4)
In [27]: np.allclose(C,E)
Out[27]: True

Upvotes: 2

Related Questions