Search code examples
pythonnumpymathmatrixtensor

What is this einsum operation doing? E = np.einsum('ij,kl->il', A,B)


given

A = np.array([[1,2],[3,4],[5,6],[7,8]])
B = np.array([[9,10,11],[12,13,14]])

matrix multiplication would be if I did

C = np.einsum('ij,jk->ik', A, B)

but if I don't multiply j in the input instead using k ...

E = np.einsum('ij,kl->il', A, B)

what am I effectively summing across? Is there a way to think about this intuitively? I'm confused because the dimensions end up the same as if it were matrix multiplication.

I've tried playing around with different numbers in the matrices A and B to get a feel for it but I'm wondering if someone can break it down for me so I can understand what is going on in this example.


Solution

  • This sums A over axis=1, B over axis=0 and computes the outer product:

    np.einsum('ij,kl->il', A, B)
    
    np.outer(A.sum(1), B.sum(0))
    
    array([[ 63,  69,  75],
           [147, 161, 175],
           [231, 253, 275],
           [315, 345, 375]])
    

    You could check the individual einsum results:

    # sum over 1
    np.einsum('ij->i', A)
    # array([ 3,  7, 11, 15])
    
    # sum over 0
    np.einsum('kl->l', B)
    # array([21, 23, 25])
    
    # outer product
    np.einsum('i,l->il', A.sum(1), B.sum(0))
    # array([[ 63,  69,  75],
    #        [147, 161, 175],
    #        [231, 253, 275],
    #        [315, 345, 375]])