Numpy einsum compute outer product along axis

I have two numpy arrays that contain compatible matrices and want to compute the element wise outer product of using numpy.einsum. The shapes of the arrays would be:

A1 = (i,j,k) A2 = (i,k,j) 

Therefore the arrays contain i matrices of shape (k,j) and (j,k) respectively.

So given A1 would contain the matrices A,B,C and A2 would contain matrices D,E,F, the result would be:

A3 = (A(x)D,B(x)E,C(x)F) 

With (x) being the outer product operator.

This would yield to my understanding based on this answer an array A3 of the following shape:

A3 = (i,j*k,j*k) 

So far I have tried:

np.einsum("ijk, ilm -> ijklm", A1, A2) 

But the resulting shapes do not fit correctly.

As a sanity check I am testing for this:

A = np.asarray(([1,2],[3,4])) B = np.asarray(([5,6],[7,8])) AB_outer = np.outer(A,B) A_vec = np.asarray((A,A)) B_vec = np.asarray((B,B)) # this line is not correct AB_vec = np.einsum("ijk, ilm -> ijklm", A_vec,B_vec) np.testing.assert_array_equal(AB_outer, AB_vec[0]) 

This currently throws an assertion error as my einsum notation is not correct. I am also open to any suggestions that can solve this and are faster or equally fast as nymphs einsum.

2 Answers

We can extend dims and let broadcasting do the job for us -

(A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3) 

Sample run -

In [46]: A1 = np.random.rand(3,4,4) ...: A2 = np.random.rand(3,4,4) In [47]: out = (A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3) In [48]: np.allclose(np.multiply.outer(A1[0],A2[0]), out[0]) Out[48]: True In [49]: np.allclose(np.multiply.outer(A1[1],A2[1]), out[1]) Out[49]: True In [50]: np.allclose(np.multiply.outer(A1[2],A2[2]), out[2]) Out[50]: True 

The equivalent with np.einsum would be -

np.einsum('ijk,ilm->ijklm',A1,A2) 

You can compute the result running:

result = np.einsum('ijk,ikl->ijl', A1, A2) 

I checked the above code on the following test data:

A = np.arange(1, 13).reshape(3, -1) B = np.arange(2, 14).reshape(3, -1) C = np.arange(3, 15).reshape(3, -1) D = np.arange(1, 13).reshape(4, -1) E = np.arange(2, 14).reshape(4, -1) F = np.arange(3, 15).reshape(4, -1) A1 = np.array([A, B, C]) A2 = np.array([D, E, F]) 

The result is:

array([[[ 70, 80, 90], [158, 184, 210], [246, 288, 330]], [[106, 120, 134], [210, 240, 270], [314, 360, 406]], [[150, 168, 186], [270, 304, 338], [390, 440, 490]]]) 

Now compute 3 "partial results":

res_1 = A @ D res_2 = B @ E res_3 = C @ F 

and check that they are just the same as consecutive sections of the result.