Numpy einsum compute outer product along axis

Question

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.

What do you mean by an outer product of matrices? An outer product is typically an operation on vectors. — user2357112, Aug 19 '20 at 10:35
@user2357112supportsMonica I am referring to the definition as described in https://math.stackexchange.com/questions/973559/outer-product-of-two-matrices. — T A, Aug 19 '20 at 10:39
@Divakar I just finished testing; works out great, thanks a lot! — T A, Aug 19 '20 at 16:57

score 4 · Accepted Answer · answered Aug 19 '20 at 10:46

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)

score 2 · Answer 2 · answered Aug 19 '20 at 18:07

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.

Numpy einsum compute outer product along axis

2 Answers2

Linked